diff options
author | komels <komels@yandex-team.ru> | 2022-04-14 13:10:53 +0300 |
---|---|---|
committer | komels <komels@yandex-team.ru> | 2022-04-14 13:10:53 +0300 |
commit | 21c9b0e6b039e9765eb414c406c2b86e8cea6850 (patch) | |
tree | f40ebc18ff8958dfbd189954ad024043ca983ea5 /kikimr | |
parent | 9a4effa852abe489707139c2b260dccc6f4f9aa9 (diff) | |
download | ydb-21c9b0e6b039e9765eb414c406c2b86e8cea6850.tar.gz |
Final part on compatibility layer: LOGBROKER-7215
ref:777c67aadbf705d19034a09a792b2df61ba53697
Diffstat (limited to 'kikimr')
135 files changed, 25662 insertions, 0 deletions
diff --git a/kikimr/.gitignore b/kikimr/.gitignore new file mode 100644 index 0000000000..504ed2bc3b --- /dev/null +++ b/kikimr/.gitignore @@ -0,0 +1,17 @@ +*.pyc +*.so +.om.cache +om.conf.local +local.cmake +local.ymake +/.gcc-compiler +/.gxx-compiler +.svn +*.log +test-results +*.a +*_pb2.py +*.pkg +*.so.1 +kikimr/driver/kikimr +.idea diff --git a/kikimr/.kikimr.root b/kikimr/.kikimr.root new file mode 100644 index 0000000000..0fe6c75b88 --- /dev/null +++ b/kikimr/.kikimr.root @@ -0,0 +1 @@ +KiKiMR project root
\ No newline at end of file diff --git a/kikimr/README.md b/kikimr/README.md new file mode 100644 index 0000000000..a0c453aa1d --- /dev/null +++ b/kikimr/README.md @@ -0,0 +1,40 @@ +# Yandex Database (YDB aka KiKiMR) and Related Projects + +See <https://wiki.yandex-team.ru/kikimr> for details. + +> For client SDK, examples, API description go to the `public` directory. + +## YDB Directory Structure + +`blockstore` -- a contract for Blockstore project (messages and their ids); +`core` -- YDB implementation; +`docs` -- documentation as a code; +`driver` -- YDB binary for Yandex (has some internal deps like Unified Agent, TVM, etc); +`driver_lib` -- a common part for `driver` and `driver_oss`; +`driver_oss` -- YDB binary for open source; +`filestore` -- a contract for Filestore project (messages and their ids); +`kikhouse_new` -- ClickHouse over YDB; +`library` -- standalone libraries used by YDB; +`mvp` -- a proxy to YDB, implements UI calls, etc; +`papers` -- papers about YDB; +`persqueue` -- persistent queue (aka Logbroker) implementation; +`public` -- SDKs, APIs and tools for YDB end users; +`services` -- grpc services for YDB public API; +`sqs` -- Amazon SQS (Simple Queue Service) API implementation on top of YDB. It is provided in Yandex.Cloud as +YMQ (Yandex Message Queue); +`streaming` -- YQL Streams implementation (as a part of Yandex Query service); +`ydbcp` -- YDB Control Plane; +`yf` -- Yandex Functions serverless computing and job execution service; +`yq` -- Yandex Query implementation; + + +## Rest +ci +deployment +juggler +production +scripts +serverless_proxy +testing +tools + diff --git a/kikimr/a.yaml b/kikimr/a.yaml new file mode 100644 index 0000000000..41209c5c88 --- /dev/null +++ b/kikimr/a.yaml @@ -0,0 +1,488 @@ +service: kikimr +title: CI for YDB +ci: + secret: sec-01ekjvt2kbg4ag64z9z03jt0hp + runtime: + sandbox-owner: STORAGE-DEV + actions: + run_sdk_java_tests: + title: Run Java tests using Maven + flow: run_java_tests + triggers: + - on: pr + filters: + - sub-paths: [ public/sdk/java/** ] + release-title-source: flow + releases: + run_periodic_export_of_data: + title: Run periodic export of data to monitoring + flow: run_periodic_export_of_data + auto: + conditions: + - schedule: + time: 9:00 - 12:00 MSK + since-last-release: 24h + stages: + package: + title: Run periodic export of data to monitoring + sync_python_sdk_from_git: + title: Sync from git + flow: sync_python_sdk_from_git + auto: + conditions: + - schedule: + time: 9:00 - 12:00 MSK + since-last-release: 24h + stages: + package: + title: Sync from git + publish_ydb_internal: + title: Build YDB Python SDK (Yandex Internal) + flow: publish_python_sdk_pypi_internal + auto: true + flow-vars: + package-path: kikimr/public/sdk/packages/python/pkg.json + stages: + package: + title: Build Python SDK + publish_ydb_persqueue_internal: + title: Build YDB Persqueue Python SDK (Yandex Internal) + flow: publish_python_sdk_pypi_internal + auto: true + flow-vars: + package-path: kikimr/public/sdk/packages/ydb-persqueue/pkg.json + stages: + package: + title: Build Python SDK + ydb_server_recipe_binary: + title: YDB Recipe Binary Release + description: | + Релиз бинарника YDB сервера в Аркадийный рецептах. + + По умолчанию, ресурс с бинарником собирается автоматически из HEAD trunk. + После сборки ресурсов создается PR и мержится после успешного прохождения тестов. + Чтобы переключить рецепт на новую мажорную версию, нужно обновить ветку в [a.yaml](http://a.yandex-team.ru/arc_vcs/kikimr/a.yaml), параметр `checkout-arcadia-from-url`. + + Автоматические релизы по понедельникам в первой половине дня. + + flow: recipe_tool + flow-vars: + build-project: kikimr/driver/kikimr + prebuilt-ya-make-path: kikimr/public/tools/package/stable + checkout-arcadia-from-url: "arcadia-arc:/#trunk" + auto: + conditions: + - schedule: + days: MON + time: 9:00 - 12:00 MSK + since-last-release: 24h + stages: + package: + title: Build binaries + lbk_recipe_tool: + title: LOGBROKER Recipe Tool Release + description: | + Релиз программы, запускающей сервер Аркадийном рецепте LOGBROKER. + + По умолчанию, рецепт обновляеся автоматически из trunk. + Релизы программы утром в понедельник. + flow: recipe_tool + flow-vars: + build-project: kikimr/public/tools/lbk_recipe/bin/lbk_recipe + prebuilt-ya-make-path: kikimr/public/tools/lbk_recipe/prebuilt + checkout-arcadia-from-url: "arcadia-arc:/#trunk" + auto: + conditions: + - schedule: + days: MON + time: 9:00 - 12:00 MSK + since-last-release: 24h + stages: + package: + title: Build binaries + sqs_recipe_tool: + title: SQS Recipe Tool Release + description: | + Релиз программы, запускающей сервер Аркадийном рецепте SQS. + + По умолчанию, рецепт обновляеся автоматически из trunk. + Релизы программы утром в понедельник. + flow: recipe_tool + flow-vars: + build-project: kikimr/public/tools/sqs_recipe/bin/sqs_recipe + prebuilt-ya-make-path: kikimr/public/tools/sqs_recipe/prebuilt + checkout-arcadia-from-url: "arcadia-arc:/#trunk" + auto: + conditions: + - schedule: + days: MON + time: 9:00 - 12:00 MSK + since-last-release: 24h + stages: + package: + title: Build binaries + ydb_recipe_tool: + title: YDB Recipe Tool Release + flow: recipe_tool + description: | + Релиз программы, запускающей сервер Аркадийном рецепте. + + По умолчанию, рецепт обновляеся автоматически из trunk. + Подробнее в [a.yaml](http://a.yandex-team.ru/arc_vcs/kikimr/a.yaml) + flow-vars: + build-project: kikimr/public/tools/ydb_recipe/bin/ydb_recipe + prebuilt-ya-make-path: kikimr/public/tools/ydb_recipe/prebuilt + checkout-arcadia-from-url: "arcadia-arc:/#trunk" + auto: + conditions: + - schedule: + days: MON + time: 9:00 - 12:00 MSK + since-last-release: 24h + stages: + package: + title: Build binaries + ydbcp_release: + title: YDB Control Plane Release + description: | + Релиз YDB Control Plane + flow: ydbcp-release-flow + start-version: 50 + branches: + pattern: releases/ydb/ydbcp/stable-${version} + auto-create: true + stages: + - id: test-ydbcp + title: Run ydbcp tests + - id: build-ydbcp + title: Build ydbcp + flows: + run_java_tests: + jobs: + run_test: + title: Build and mvn install kikimr/public/sdk/java + task: common/misc/run_command + requirements: + cores: 16 + ram: 32 GB + sandbox: + client_tags: GENERIC & LINUX & SSD + container_resource: 2309143241 # Container with Java 8, Java 11 and Maven + dns: dns64 + input: + config: + cmd_line: | + set -ex + cd $ARCADIA_PATH + export PATH=$ARCADIA_PATH:$PATH + export JAVA_HOME=/usr/local/java8 + mkdir -p ~/.m2 + cp ci/tasklet/registry/common/misc/run_command/maven/settings.xml ~/.m2/ + cd kikimr/public/sdk/java + ya make + ./gen-proto.sh + mvn -B -Pwith-examples install + secret_environment_variables: + - key: YA_TOKEN + secret_spec: + uuid: sec-01ekjvt2kbg4ag64z9z03jt0hp + key: ci.token + arc_mount_config: + enabled: true + run_periodic_export_of_data: + jobs: + run_test: + title: run_periodic_export_of_data + task: common/misc/run_command + requirements: + sandbox: + client_tags: GENERIC & LINUX + dns: dns64 + input: + config: + secret_environment_variables: + - key: YC_OAUTH_TOKEN + secret_spec: + uuid: sec-01fzd4egdm6zk7bfxd3bnf3fnw + key: token + - key: YA_TOKEN + secret_spec: + uuid: sec-01ekjvt2kbg4ag64z9z03jt0hp + key: ci.token + - key: YT_TOKEN + secret_spec: + uuid: sec-01fzd4egdm6zk7bfxd3bnf3fnw + key: yt + arc_mount_config: + enabled: true + cmd_line: | + set -xe + curl https://s3.mds.yandex.net/mcdev/ycp/install.sh | bash + export PATH=/home/sandbox/ycp/bin:$PATH + ycp init + cp $ARCADIA_PATH/kikimr/ydbcp/abc_database/config.yaml /home/sandbox/.config/ycp/config.yaml + sed -i "s/TOKEN_TO_REPLACE/"$YC_OAUTH_TOKEN"/g" /home/sandbox/.config/ycp/config.yaml + ycp --format json --profile internal team integration abc resolve -r - <<<"abc_slug: kikimr" + cd $ARCADIA_PATH/kikimr/ydbcp/abc_database/ + $ARCADIA_PATH/ya make -r + ./collect_info + cat result.txt + cat result.txt | YT_PROXY=hahn $ARCADIA_PATH/ya tool yt write --table //home/solomon/service_provider_alerts/service_provider_exports/ydb --format json + run_build_from_git: + jobs: + run_test: + title: Build python sdk and then run tests + task: common/misc/run_command + requirements: + cores: 56 + ram: 256 GB + sandbox: + client_tags: GENERIC & LINUX & SSD + dns: dns64 + input: + config: + result_resources: + - path: package + description: YDBD TAR + type: OTHER_RESOURCE + compression_type: tgz + ci_badge: true + attributes: + released: stable + ttl: "14" + cmd_line: | + git clone https://github.yandex-team.ru/arcadia-devtools/ydb_oss.git + mkdir $RESULT_RESOURCES_PATH/package + cd ydb_oss + ./ya make -r -k ydb/apps/ydbd/ + cp ydb/apps/ydbd/ydbd $RESULT_RESOURCES_PATH/package/ + cp ydb/apps/ydbd/*.so $RESULT_RESOURCES_PATH/package/ + publish_python_sdk_pypi_internal: + title: Build YDB Python SDK (Yandex Internal) + jobs: + package: + title: Build YDB Python SDK (Yandex Internal) + description: Runs ya package in sandbox + task: common/arcadia/ya_package + input: + packages: ${flow-vars.package-path} + build_type: release + package_type: wheel + publish_package: True + publish_to: "" + wheel_access_key_token: robot_ydb_pypi_access_token + wheel_secret_key_token: robot_ydb_pypi_key_token + wheel_upload_repo: https://pypi.yandex-team.ru/simple + build_system: ya + sandbox_container: 1329390397 + stage: package + ydbcp-release-flow: + title: Release YDBCP + jobs: + start-release: + title: Start release + task: dummy + stage: test-ydbcp + run-ydbcp-tests: + title: Run ydbcp tests + description: Run ydbcp tests + task: common/arcadia/ya_make + stage: test-ydbcp + needs: start-release + requirements: + sandbox: + client_tags: GENERIC & LINUX & SSD + cores: 56 + ram: 256GB + input: &run-tests-on-distbuild-input + build_system: ya_force + test: True + keep_on: True + checkout: False + checkout_mode: manual + output_only_tests: True + test_log_level: debug + collect_test_cores: False + do_not_restart: True + strip_binaries: True + check_return_code: True + fail_on_any_error: True + junit_report: True + kill_timeout: 14400 + build_output_ttl: 1 + build_output_html_ttl: 14 + ya_timeout: 13500 + test_threads: 50 + # distbuild_pool: "//man/users/ydb" + # env_vars: "YA_TOKEN='$(vault:value:STORAGE-DEV:ARC_TOKEN)'" + targets: "kikimr/ydbcp" + validate-test-failures: + title: Validate test fails + task: dummy + manual: + enabled: true + prompt: "Ignore test failures?" + needs: start-release + run-create-tag: + title: Make YDBCP release tag + description: "Mount ARC and create the tag" + task: common/misc/run_command + requirements: + cores: 2 + ram: 16G + sandbox: + client_tags: GENERIC & LINUX + needs-type: any + needs: + - run-ydbcp-tests + - validate-test-failures + stage: build-ydbcp + input: + config: + arc_mount_config: + enabled: true + revision_hash: trunk + environment_variables: + - key: ARC_TAG + value: "tags/releases/ydb/ydbcp/stable-${context.version_info.major}-${not_null(context.version_info.minor, `0`)}" + - key: ARC_COMMIT + value: "${context.target_revision.hash}" + secret_environment_variables: + - key: YA_TOKEN + secret_spec: + uuid: sec-01ekjvt2kbg4ag64z9z03jt0hp + key: ci.token + cmd_line: | + pwd + $ARC_BIN status + $ARC_BIN push $ARC_COMMIT:$ARC_TAG + build-ydbcp: + title: Run build YDBCP packages + description: Run build YDB packages + task: common/arcadia/ya_package_2 + needs: run-create-tag + stage: build-ydbcp + requirements: + sandbox: + client_tags: GENERIC & LINUX & SSD + cores: 56 + ram: 256GB + input: + build_system: ya_force + checkout_arcadia_from_url: "arcadia-arc:/#tags/releases/ydb/ydbcp/stable-${context.version_info.major}-${not_null(context.version_info.minor, `0`)}" + key_user: robot-simcity + publish_package: True + multiple_publish: True + publish_to: search + multiple_publish_to: search + build_logs_ttl: 1 + package_ttl: 1 + package_type: debian + build_output_ttl: 1 + packages: "kikimr/ydbcp/packages/ydb-ydbcp-bin/pkg.json;kikimr/ydbcp/packages/ydb-ydbcp-init/pkg.json;kikimr/ydbcp/packages/ydb-ydbcp-prestable-meta-cfg/pkg.json;kikimr/ydbcp/packages/ydb-ydbcp-prod-meta-cfg/pkg.json;" + sync_python_sdk_from_git: + title: Sync from git + jobs: + sync_ydb_python_sdk_from_git: + title: Make PR with update + description: Runs + task: common/misc/run_command + requirements: + cores: 2 + ram: 16G + sandbox: + client_tags: GENERIC & LINUX + dns: dns64 + input: + config: + arc_mount_config: + enabled: true + revision_hash: trunk + environment_variables: + - key: ARC_BRANCH + value: "users/robot-kikimr-dev/execute-sync-from-github-${context.version_info.major}" + secret_environment_variables: + - key: YA_TOKEN + secret_spec: + uuid: sec-01ekjvt2kbg4ag64z9z03jt0hp + key: ci.token + cmd_line: | + pwd + $ARC_BIN status + $ARC_BIN info --json + $ARC_BIN checkout -b execute-sync-from-github + git clone https://github.com/yandex-cloud/ydb-python-sdk ~/tmp + rm -rf ydb/public/sdk/python/* + cp -r ~/tmp/* ydb/public/sdk/python/ + python3 kikimr/public/sdk/python/make_ya_make.py + $ARC_BIN add -f ydb/public/sdk/python + $ARC_BIN commit -am "Release YDB recipe binaries" + $ARC_BIN push --set-upstream $ARC_BRANCH + $ARC_BIN pr create --publish -m 'Sync YDB SDK from github' --no-edit + recipe_tool: + title: Recipe tool binary + jobs: + start-build: + title: Build applications + task: dummy + stage: package + build-project: + title: Build Recipe Tool + description: Runs sandbox task + task: common/arcadia/build_arcadia_project_for_all + needs: + - start-build + input: + project: ${flow-vars.build-project} + checkout_arcadia_from_url: ${flow-vars.checkout-arcadia-from-url} + platforms_list: + - linux + - darwin + strip_binaries: True + definition_flags: -DUSE_YDB_TRUNK_RECIPE_TOOLS=1 + recipe_pr: + title: Make PR with resources + description: Runs + task: common/misc/run_command + requirements: + cores: 2 + ram: 16G + sandbox: + client_tags: GENERIC & LINUX + needs: + - start-build + - build-project + input: + config: + arc_mount_config: + enabled: true + revision_hash: trunk + environment_variables: + - key: YDB_LINUX + value: ${to_string((tasks.build-project.resources[?type == 'ARCADIA_PROJECT_TGZ' && attributes.platform == 'linux'])[0].id)} + - key: YDB_DARWIN + value: ${to_string((tasks.build-project.resources[?type == 'ARCADIA_PROJECT_TGZ' && attributes.platform == 'darwin'])[0].id)} + - key: RECIPE_TOOL_PATH + value: ${flow-vars.prebuilt-ya-make-path} + secret_environment_variables: + - key: YA_TOKEN + secret_spec: + uuid: sec-01ekjvt2kbg4ag64z9z03jt0hp + key: ci.token + cmd_line: | + pwd + $ARC_BIN status + $ARC_BIN info --json + $ARC_BIN checkout -b release-recipe-binaries-$YDB_LINUX-$YDB_DARWIN + rm $RECIPE_TOOL_PATH/ya.make + cp $RECIPE_TOOL_PATH/ya.make-template $RECIPE_TOOL_PATH/ya.make + sed -i 's/YDB_LINUX/'$YDB_LINUX'/g' $RECIPE_TOOL_PATH/ya.make + sed -i 's/YDB_DARWIN/'$YDB_DARWIN'/g' $RECIPE_TOOL_PATH/ya.make + $ARC_BIN add $RECIPE_TOOL_PATH/ya.make + $ARC_BIN commit -am "Release YDB recipe binaries" + $ARC_BIN push --set-upstream users/robot-kikimr-dev/release-recipe-binaries-$YDB_LINUX-$YDB_DARWIN + $ARC_BIN pr create --publish -m 'Release YDB recipe binaries' --no-edit + autocheck: + strong: true + fast-targets: + kikimr diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/CMakeLists.txt b/kikimr/persqueue/sdk/deprecated/cpp/v2/CMakeLists.txt new file mode 100644 index 0000000000..44cb2496ed --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/CMakeLists.txt @@ -0,0 +1,60 @@ + +# This file was gererated by the build system used internally in the Yandex monorepo. +# Only simple modifications are allowed (adding source-files to targets, adding simple properties +# like target_include_directories). These modifications will be ported to original +# ya.make files by maintainers. Any complex modifications which can't be ported back to the +# original buildsystem will not be accepted. + + + +add_library(deprecated-cpp-v2) +target_link_libraries(deprecated-cpp-v2 PUBLIC + contrib-libs-cxxsupp + yutil + contrib-libs-grpc + contrib-libs-protobuf + cpp-client-iam + cpp-client-ydb_persqueue + api-grpc-yndx + api-protos-yndx + yndx-persqueue-read_batch_converter + library-cpp-blockcodecs + cpp-containers-disjoint_interval_tree + cpp-containers-intrusive_rb_tree + cpp-grpc-common + cpp-http-simple + library-cpp-json + library-cpp-logger + cpp-streams-lzop + cpp-streams-zstd + cpp-string_utils-quote + cpp-threading-future + cpp-tvmauth-client + library-persqueue-topic_parser_public + cpp-client-resources +) +target_sources(deprecated-cpp-v2 PRIVATE + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/channel.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compressing_producer.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/decompressing_consumer.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iconsumer_p.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/interface_common.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iprocessor_p.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iproducer_p.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_consumer.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_producer.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/persqueue.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/processor.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/producer.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/logger.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/types.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/queue_pool.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_consumer.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_producer.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/scheduler.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/consumer.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/credentials_provider.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/validate_grpc_metadata.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compat_producer.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/ydb_sdk_consumer.cpp +) diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/actor_interface.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/actor_interface.h new file mode 100644 index 0000000000..623a293dc3 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/actor_interface.h @@ -0,0 +1,60 @@ +#pragma once +#include "fake_actor.h" + +#include <library/cpp/actors/core/actor.h> +#include <library/cpp/actors/core/actorid.h> +#include <library/cpp/actors/core/actorsystem.h> +#include <library/cpp/actors/core/events.h> + +namespace NPersQueue { + +struct IActorInterface { + virtual ~IActorInterface() = default; + + // Sender of all messages from PQLib objects. + virtual NActors::TActorId GetActorID() const noexcept = 0; +}; + +class TActorHolder : public IActorInterface +{ +protected: + TActorHolder(NActors::TActorSystem* actorSystem, const NActors::TActorId& parentActorID) + : ActorSystem(actorSystem) + , ParentActorID(parentActorID) + { + ActorID = ActorSystem->Register(new TFakeActor()); + } + + ~TActorHolder() { + ActorSystem->Send(new NActors::IEventHandle(ActorID, ActorID, new NActors::TEvents::TEvPoisonPill())); + } + + template <class TResponseEvent> + void Subscribe(ui64 requestId, NThreading::TFuture<typename TResponseEvent::TResponseType>&& future) { + auto handler = [ + requestId = requestId, + actorSystem = ActorSystem, + actorID = ActorID, + parentActorID = ParentActorID + ](const NThreading::TFuture<typename TResponseEvent::TResponseType>& future) { + actorSystem->Send(new NActors::IEventHandle( + parentActorID, + actorID, + new TResponseEvent(typename TResponseEvent::TResponseType(future.GetValue()), requestId) + )); + }; + future.Subscribe(handler); + } + +public: + NActors::TActorId GetActorID() const noexcept override { + return ActorID; + } + +protected: + NActors::TActorSystem* ActorSystem; + NActors::TActorId ParentActorID; + NActors::TActorId ActorID; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/actor_wrappers.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/actor_wrappers.h new file mode 100644 index 0000000000..ce2f7682c1 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/actor_wrappers.h @@ -0,0 +1,108 @@ +#pragma once +#include "actor_interface.h" +#include "responses.h" + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iproducer.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iconsumer.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iprocessor.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> + +#include <util/datetime/base.h> + +namespace NPersQueue { + +template <ui32 TObjectIsDeadEventId, ui32 TCreateResponseEventId, ui32 TCommitResponseEventId> +class TProducerActorWrapper: public TActorHolder { +public: + using TObjectIsDeadEvent = TPQLibResponseEvent<TError, TObjectIsDeadEventId>; + using TCreateResponseEvent = TPQLibResponseEvent<TProducerCreateResponse, TCreateResponseEventId>; + using TCommitResponseEvent = TPQLibResponseEvent<TProducerCommitResponse, TCommitResponseEventId>; + +public: + TProducerActorWrapper(NActors::TActorSystem* actorSystem, const NActors::TActorId& parentActorID, THolder<IProducer> producer) + : TActorHolder(actorSystem, parentActorID) + , Producer(std::move(producer)) + { + } + + void Start(TInstant deadline, ui64 requestId = 0, ui64 isDeadRequestId = 0) noexcept { + Subscribe<TCreateResponseEvent>(requestId, Producer->Start(deadline)); + Subscribe<TObjectIsDeadEvent>(isDeadRequestId, Producer->IsDead()); + } + + void Start(TDuration timeout, ui64 requestId = 0, ui64 isDeadRequestId = 0) noexcept { + Start(TInstant::Now() + timeout, requestId, isDeadRequestId); + } + + void Write(TProducerSeqNo seqNo, TData data, ui64 requestId = 0) noexcept { + Subscribe<TCommitResponseEvent>(requestId, Producer->Write(seqNo, std::move(data))); + } + + void Write(TData data, ui64 requestId = 0) noexcept { + Subscribe<TCommitResponseEvent>(requestId, Producer->Write(std::move(data))); + } + +private: + THolder<IProducer> Producer; +}; + +template <ui32 TObjectIsDeadEventId, ui32 TCreateResponseEventId, ui32 TGetMessageEventId> +class TConsumerActorWrapper: public TActorHolder { +public: + using TObjectIsDeadEvent = TPQLibResponseEvent<TError, TObjectIsDeadEventId>; + using TCreateResponseEvent = TPQLibResponseEvent<TConsumerCreateResponse, TCreateResponseEventId>; + using TGetMessageEvent = TPQLibResponseEvent<TConsumerMessage, TGetMessageEventId>; + +public: + TConsumerActorWrapper(NActors::TActorSystem* actorSystem, const NActors::TActorId& parentActorID, THolder<IConsumer> consumer) + : TActorHolder(actorSystem, parentActorID) + , Consumer(std::move(consumer)) + { + } + + void Start(TInstant deadline, ui64 requestId = 0, ui64 isDeadRequestId = 0) noexcept { + Subscribe<TCreateResponseEvent>(requestId, Consumer->Start(deadline)); + Subscribe<TObjectIsDeadEvent>(isDeadRequestId, Consumer->IsDead()); + } + + void Start(TDuration timeout, ui64 requestId = 0, ui64 isDeadRequestId = 0) noexcept { + Start(TInstant::Now() + timeout, requestId, isDeadRequestId); + } + + void GetNextMessage(ui64 requestId = 0) noexcept { + Subscribe<TGetMessageEvent>(requestId, Consumer->GetNextMessage()); + } + + void Commit(const TVector<ui64>& cookies) noexcept { + Consumer->Commit(cookies); // no future in PQLib API + } + + void RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept { + Consumer->RequestPartitionStatus(topic, partition, generation); // no future in PQLib API + } + +private: + THolder<IConsumer> Consumer; +}; + +template <ui32 TOriginDataEventId> +class TProcessorActorWrapper: public TActorHolder { +public: + using TOriginDataEvent = TPQLibResponseEvent<TOriginData, TOriginDataEventId>; + +public: + TProcessorActorWrapper(NActors::TActorSystem* actorSystem, const NActors::TActorId& parentActorID, THolder<IProcessor> processor) + : TActorHolder(actorSystem, parentActorID) + , Processor(std::move(processor)) + { + } + + void GetNextData(ui64 requestId = 0) noexcept { + Subscribe<TOriginDataEvent>(requestId, Processor->GetNextData()); + } + +private: + THolder<IProcessor> Processor; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/actor_wrappers_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/actor_wrappers_ut.cpp new file mode 100644 index 0000000000..1e08597ec3 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/actor_wrappers_ut.cpp @@ -0,0 +1,119 @@ +#include "actor_wrappers.h" +#include "responses.h" + +#include <ydb/core/testlib/basics/runtime.h> +#include <ydb/core/testlib/basics/appdata.h> + +#include <library/cpp/actors/core/events.h> +#include <library/cpp/testing/gmock_in_unittest/gmock.h> +#include <library/cpp/testing/unittest/registar.h> + +#include <util/system/thread.h> + +using namespace testing; + +namespace NPersQueue { + +namespace TEvPQLibTests { + enum EEv { + EvPQLibObjectIsDead = EventSpaceBegin(NActors::TEvents::ES_PRIVATE), + EvPQLibProducerCreateResponse, + EvPQLibProducerCommitResponse, + EvEnd + }; +} // TEvPQLibTests + +class TMockProducer: public IProducer { +public: + NThreading::TFuture<TProducerCreateResponse> Start(TInstant) noexcept override { + return NThreading::MakeFuture<TProducerCreateResponse>(MockStart()); + } + + NThreading::TFuture<TProducerCommitResponse> Write(TProducerSeqNo, TData data) noexcept override { + return Write(std::move(data)); + } + + NThreading::TFuture<TProducerCommitResponse> Write(TData data) noexcept override { + return NThreading::MakeFuture<TProducerCommitResponse>(MockWrite(data)); + } + + NThreading::TFuture<TError> IsDead() noexcept override { + return MockIsDead(); + } + + MOCK_METHOD(TProducerCreateResponse, MockStart, (), ()); + MOCK_METHOD(NThreading::TFuture<TError>, MockIsDead, (), ()); + MOCK_METHOD(TProducerCommitResponse, MockWrite, (TData), ()); +}; + +Y_UNIT_TEST_SUITE(TProducerActorWrapperTest) { + Y_UNIT_TEST(PassesEvents) { + NActors::TTestActorRuntime runtime; + runtime.Initialize(NKikimr::TAppPrepare().Unwrap()); + auto edgeActorId = runtime.AllocateEdgeActor(); + THolder<TMockProducer> producerHolder = MakeHolder<TMockProducer>(); + TMockProducer& producer = *producerHolder; + + // expectations + TWriteResponse startResponse; + startResponse.MutableInit()->SetMaxSeqNo(42); + EXPECT_CALL(producer, MockStart()) + .WillOnce(Return(TProducerCreateResponse(std::move(startResponse)))); + + NThreading::TPromise<TError> deadPromise = NThreading::NewPromise<TError>(); + EXPECT_CALL(producer, MockIsDead()) + .WillOnce(Return(deadPromise.GetFuture())); + + TWriteResponse writeResponse; + writeResponse.MutableAck()->SetSeqNo(100); + TData data("data"); + EXPECT_CALL(producer, MockWrite(data)) + .WillOnce(Return(TProducerCommitResponse(100, data, std::move(writeResponse)))); + + // wrapper creation + using TProducerType = TProducerActorWrapper< + TEvPQLibTests::EvPQLibObjectIsDead, + TEvPQLibTests::EvPQLibProducerCreateResponse, + TEvPQLibTests::EvPQLibProducerCommitResponse + >; + + THolder<TProducerType> wrapper = MakeHolder<TProducerType>(runtime.GetAnyNodeActorSystem(), edgeActorId, std::move(producerHolder)); + + // checks + auto actorId = wrapper->GetActorID(); + UNIT_ASSERT(actorId); + + { + wrapper->Start(TInstant::Max()); + TAutoPtr<NActors::IEventHandle> handle; + auto* startEvent = runtime.GrabEdgeEvent<TProducerType::TCreateResponseEvent>(handle); + UNIT_ASSERT(startEvent != nullptr); + UNIT_ASSERT_VALUES_EQUAL(handle->Sender, actorId); + UNIT_ASSERT_VALUES_EQUAL(startEvent->Response.Response.GetInit().GetMaxSeqNo(), 42); + } + + { + wrapper->Write(data); + TAutoPtr<NActors::IEventHandle> handle; + auto* writeEvent = runtime.GrabEdgeEvent<TProducerType::TCommitResponseEvent>(handle); + UNIT_ASSERT(writeEvent != nullptr); + UNIT_ASSERT_VALUES_EQUAL(handle->Sender, actorId); + UNIT_ASSERT_VALUES_EQUAL(writeEvent->Response.Response.GetAck().GetSeqNo(), 100); + UNIT_ASSERT_EQUAL(writeEvent->Response.Data, data); + UNIT_ASSERT_VALUES_EQUAL(writeEvent->Response.SeqNo, 100); + } + + { + TError err; + err.SetDescription("trololo"); + deadPromise.SetValue(err); + + TAutoPtr<NActors::IEventHandle> handle; + auto* deadEvent = runtime.GrabEdgeEvent<TProducerType::TObjectIsDeadEvent>(handle); + UNIT_ASSERT(deadEvent != nullptr); + UNIT_ASSERT_VALUES_EQUAL(handle->Sender, actorId); + UNIT_ASSERT_STRINGS_EQUAL(deadEvent->Response.GetDescription(), "trololo"); + } + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/fake_actor.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/fake_actor.h new file mode 100644 index 0000000000..630e8a35f8 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/fake_actor.h @@ -0,0 +1,24 @@ +#pragma once +#include <library/cpp/actors/core/actor.h> +#include <library/cpp/actors/core/events.h> +#include <library/cpp/actors/core/hfunc.h> + +namespace NPersQueue { + +// Fake object to hold actor id +class TFakeActor: public NActors::TActor<TFakeActor> { +public: + TFakeActor() + : TActor<TFakeActor>(&TFakeActor::StateFunc) + { + } + +private: + STFUNC(StateFunc) { + switch (ev->GetTypeRewrite()) { + CFunc(NActors::TEvents::TSystem::PoisonPill, Die); + } + } +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/logger.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/logger.h new file mode 100644 index 0000000000..4ad5fb0d48 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/logger.h @@ -0,0 +1,48 @@ +#pragma once +#include "responses.h" +#include "actor_interface.h" +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/logger.h> + +namespace NPersQueue { + +template <ui32 EventTypeId> +class TActorLogger + : public TActorHolder + , public ILogger +{ +public: + using TLogEvent = TEvPQLibLog<EventTypeId>; + +public: + explicit TActorLogger(NActors::TActorSystem* actorSystem, const NActors::TActorId& parentActorID, int level) + : TActorHolder(actorSystem, parentActorID) + , Level(level) + { + } + + ~TActorLogger() override = default; + + void Log(const TString& msg, const TString& sourceId, const TString& sessionId, int level) override { + if (!IsEnabled(level)) { + return; + } + ActorSystem->Send(ParentActorID, new TLogEvent( + msg, + sourceId, + sessionId, + static_cast<NActors::NLog::EPriority>(level) + )); + } + + bool IsEnabled(int level) const override { + return level <= Level; + } + +private: + int Level; +}; + +} + + + diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/persqueue.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/persqueue.cpp new file mode 100644 index 0000000000..c663c3ac04 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/persqueue.cpp @@ -0,0 +1,22 @@ +#include "actor_wrappers.h" +#include "persqueue.h" + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/impl/internals.h> + +namespace NPersQueue { + +TPQLibActorsWrapper::TPQLibActorsWrapper(NActors::TActorSystem* actorSystem, const TPQLibSettings& settings) + : PQLib(new TPQLib(settings)) + , ActorSystem(actorSystem) +{ +} + +TPQLibActorsWrapper::TPQLibActorsWrapper(NActors::TActorSystem* actorSystem, std::shared_ptr<TPQLib> pqLib) + : PQLib(std::move(pqLib)) + , ActorSystem(actorSystem) +{ +} + +TPQLibActorsWrapper::~TPQLibActorsWrapper() = default; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/persqueue.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/persqueue.h new file mode 100644 index 0000000000..7a11416f4f --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/persqueue.h @@ -0,0 +1,87 @@ +#pragma once +#include "actor_wrappers.h" + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> + +#include <library/cpp/actors/core/actorid.h> + +#include <util/generic/noncopyable.h> +#include <util/generic/ptr.h> + +namespace NActors { +class TActorSystem; +} // namespace NActors + +namespace NPersQueue { + +class TPQLibActorsWrapper : public TNonCopyable { +public: + explicit TPQLibActorsWrapper(NActors::TActorSystem* actorSystem, const TPQLibSettings& settings = TPQLibSettings()); + TPQLibActorsWrapper(NActors::TActorSystem* actorSystem, std::shared_ptr<TPQLib> pqLib); + ~TPQLibActorsWrapper(); + + // Producers creation + template <ui32 TObjectIsDeadEventId, ui32 TCreateResponseEventId, ui32 TCommitResponseEventId> + THolder<TProducerActorWrapper<TObjectIsDeadEventId, TCreateResponseEventId, TCommitResponseEventId>> CreateProducer( + const NActors::TActorId& parentActorID, + const TProducerSettings& settings, + TIntrusivePtr<ILogger> logger = nullptr, + bool deprecated = false + ) { + return new TProducerActorWrapper<TObjectIsDeadEventId, TCreateResponseEventId, TCommitResponseEventId>( + ActorSystem, + parentActorID, + PQLib->CreateProducer(settings, std::move(logger), deprecated) + ); + } + + template <ui32 TObjectIsDeadEventId, ui32 TCreateResponseEventId, ui32 TCommitResponseEventId> + THolder<TProducerActorWrapper<TObjectIsDeadEventId, TCreateResponseEventId, TCommitResponseEventId>> CreateMultiClusterProducer( + const NActors::TActorId& parentActorID, + const TMultiClusterProducerSettings& settings, + TIntrusivePtr<ILogger> logger = nullptr, + bool deprecated = false + ) { + return new TProducerActorWrapper<TObjectIsDeadEventId, TCreateResponseEventId, TCommitResponseEventId>( + ActorSystem, + parentActorID, + PQLib->CreateMultiClusterProducer(settings, std::move(logger), deprecated) + ); + } + + // Consumers creation + template <ui32 TObjectIsDeadEventId, ui32 TCreateResponseEventId, ui32 TGetMessageEventId> + THolder<TConsumerActorWrapper<TObjectIsDeadEventId, TCreateResponseEventId, TGetMessageEventId>> CreateConsumer( + const NActors::TActorId& parentActorID, + const TConsumerSettings& settings, + TIntrusivePtr<ILogger> logger = nullptr, + bool deprecated = false + ) { + return MakeHolder<TConsumerActorWrapper<TObjectIsDeadEventId, TCreateResponseEventId, TGetMessageEventId>>( + ActorSystem, + parentActorID, + PQLib->CreateConsumer(settings, std::move(logger),deprecated) + ); + } + + // Processors creation + template <ui32 TOriginDataEventId> + THolder<TProcessorActorWrapper<TOriginDataEventId>> CreateProcessor( + const NActors::TActorId& parentActorID, + const TProcessorSettings& settings, + TIntrusivePtr<ILogger> logger = nullptr, + bool deprecated = false + ) { + return new TProcessorActorWrapper<TOriginDataEventId>( + ActorSystem, + parentActorID, + PQLib->CreateProcessor(settings, std::move(logger), deprecated) + ); + } + +private: + std::shared_ptr<TPQLib> PQLib; + NActors::TActorSystem* ActorSystem; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/responses.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/responses.h new file mode 100644 index 0000000000..86e732d87c --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/actors/responses.h @@ -0,0 +1,39 @@ +#pragma once +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iprocessor.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> + +#include <library/cpp/actors/core/events.h> +#include <library/cpp/actors/core/event_local.h> + +namespace NPersQueue { + +template <class TResponseType_, ui32 EventTypeId> +struct TPQLibResponseEvent: public NActors::TEventLocal<TPQLibResponseEvent<TResponseType_, EventTypeId>, EventTypeId> { + using TResponseType = TResponseType_; + + explicit TPQLibResponseEvent(TResponseType&& response, ui64 requestId = 0) + : Response(std::move(response)) + , RequestId(requestId) + { + } + + TResponseType Response; + ui64 RequestId; +}; + +template <ui32 EventTypeId> +struct TEvPQLibLog : NActors::TEventLocal<TEvPQLibLog<EventTypeId>, EventTypeId> { + TEvPQLibLog(const TString& msg, const TString& sourceId, const TString& sessionId, NActors::NLog::EPriority level) + : Message(msg) + , SourceId(sourceId) + , SessionId(sessionId) + , Level(level) + {} + + TString Message; + TString SourceId; + TString SessionId; + NActors::NLog::EPriority Level; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/compatibility_ut/compatibility_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/compatibility_ut/compatibility_ut.cpp new file mode 100644 index 0000000000..b4d8431c8a --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/compatibility_ut/compatibility_ut.cpp @@ -0,0 +1,113 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h> + +#include <library/cpp/testing/unittest/registar.h> + +using namespace NThreading; +using namespace NKikimr; +using namespace NKikimr::NPersQueueTests; + +namespace NPersQueue { +Y_UNIT_TEST_SUITE(TCompatibilityTest) { + void AssertWriteValid(const NThreading::TFuture<TProducerCommitResponse>& respFuture) { + const TProducerCommitResponse& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TWriteResponse::kAck, "Msg: " << resp.Response); + } + + void AssertWriteFailed(const NThreading::TFuture<TProducerCommitResponse>& respFuture) { + const TProducerCommitResponse& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TWriteResponse::kError, "Msg: " << resp.Response); + } + + void AssertReadValid(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kData, "Msg: " << resp.Response); + } + + void AssertReadContinuous(const TReadResponse::TData::TMessageBatch& batch, ui64 startSeqNo) { + for (ui32 i = 0; i < batch.MessageSize(); ++i) { + ui64 actualSeqNo = batch.GetMessage(i).GetMeta().GetSeqNo(); + UNIT_ASSERT_EQUAL_C(actualSeqNo, startSeqNo + i, "Wrong seqNo: " << actualSeqNo << ", expected: " << startSeqNo + i); + } + } + + void AssertCommited(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kCommit, "Msg: " << resp.Response); + } + + void AssertReadFailed(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kError, "Msg: " << resp.Response); + } + + void AssertLock(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kLock, "Msg: " << resp.Response); + } + + TProducerSettings MakeProducerSettings(const TTestServer& testServer) { + TProducerSettings producerSettings; + producerSettings.ReconnectOnFailure = true; + producerSettings.Topic = "topic1"; + producerSettings.SourceId = "123"; + producerSettings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + producerSettings.Codec = ECodec::LZOP; + producerSettings.ReconnectionDelay = TDuration::MilliSeconds(10); + return producerSettings; + } + + TConsumerSettings MakeConsumerSettings(const TTestServer& testServer) { + TConsumerSettings consumerSettings; + consumerSettings.ClientId = "user"; + consumerSettings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + consumerSettings.Topics.push_back("topic1"); + return consumerSettings; + } + + Y_UNIT_TEST(ContinuesOperationsAfterPQLibDeath) { + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(); + + const size_t partitions = 1; + testServer.AnnoyingClient->FullInit(); + testServer.AnnoyingClient->CreateTopic("rt3.dc1--topic1", partitions); + + testServer.WaitInit("topic1"); + + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLibSettings pqLibSettings; + pqLibSettings.DefaultLogger = logger; + THolder<TPQLib> PQLib = MakeHolder<TPQLib>(pqLibSettings); + + auto producer = PQLib->CreateProducer(MakeProducerSettings(testServer), logger, true); + UNIT_ASSERT(!producer->Start().GetValueSync().Response.HasError()); + auto isProducerDead = producer->IsDead(); + UNIT_ASSERT(!isProducerDead.HasValue()); + + auto consumer = PQLib->CreateConsumer(MakeConsumerSettings(testServer), logger, true); + + PQLib = nullptr; + + UNIT_ASSERT(!consumer->Start().GetValueSync().Response.HasError()); + auto isConsumerDead = consumer->IsDead(); + UNIT_ASSERT(!isConsumerDead.HasValue()); + + auto write1 = producer->Write(1, TString("blob1")); + auto write2 = producer->Write(2, TString("blob2")); + + UNIT_ASSERT(!write1.GetValueSync().Response.HasError()); + UNIT_ASSERT(!write2.GetValueSync().Response.HasError()); + + UNIT_ASSERT(!consumer->GetNextMessage().GetValueSync().Response.HasError()); + + producer = nullptr; + consumer = nullptr; + + isProducerDead.GetValueSync(); + isConsumerDead.GetValueSync(); + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/credentials_provider.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/credentials_provider.h new file mode 100644 index 0000000000..48df7b0cdd --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/credentials_provider.h @@ -0,0 +1,63 @@ +#pragma once + +#include "logger.h" + +#include <kikimr/yndx/api/protos/persqueue.pb.h> + +#include <library/cpp/tvmauth/client/facade.h> + +namespace NPersQueue { + +const TString& GetTvmPqServiceName(); + +class ICredentialsProvider { +public: + virtual ~ICredentialsProvider() = default; + virtual void FillAuthInfo(NPersQueueCommon::TCredentials* authInfo) const = 0; +}; + + +class ITVMCredentialsForwarder : public ICredentialsProvider { +public: + virtual ~ITVMCredentialsForwarder() = default; + virtual void SetTVMServiceTicket(const TString& ticket) = 0; +}; + +std::shared_ptr<ICredentialsProvider> CreateInsecureCredentialsProvider(); + +/* + * Builder for direct usage of TVM. For Qloud environment use CreateTVMQloudCredentialsProvider + */ +std::shared_ptr<ICredentialsProvider> CreateTVMCredentialsProvider(const TString& secret, const ui32 srcClientId, const ui32 dstClientId, TIntrusivePtr<ILogger> logger = nullptr); +std::shared_ptr<ICredentialsProvider> CreateTVMCredentialsProvider(const NTvmAuth::NTvmApi::TClientSettings& settings, TIntrusivePtr<ILogger> logger = nullptr, const TString& dstAlias = GetTvmPqServiceName()); + +/* + * forward TVM service ticket with ITMVCredentialsProvider method SetTVMServiceTicket ; threadsafe + */ +std::shared_ptr<ITVMCredentialsForwarder> CreateTVMCredentialsForwarder(); + +/* + * tvmClient settings must contain LogBroker client id under an alias + */ +std::shared_ptr<ICredentialsProvider> CreateTVMCredentialsProvider(std::shared_ptr<NTvmAuth::TTvmClient> tvmClient, TIntrusivePtr<ILogger> logger = nullptr, const TString& dstAlias = GetTvmPqServiceName()); + +std::shared_ptr<ICredentialsProvider> CreateTVMQloudCredentialsProvider(const TString& srcAlias, const TString& dstAlias, TIntrusivePtr<ILogger> logger = nullptr, const TDuration refreshPeriod = TDuration::Minutes(10), ui32 port = 1); +std::shared_ptr<ICredentialsProvider> CreateTVMQloudCredentialsProvider(const ui32 srcId, const ui32 dstId, TIntrusivePtr<ILogger> logger = nullptr, const TDuration refreshPeriod = TDuration::Minutes(10), ui32 port = 1); +std::shared_ptr<ICredentialsProvider> CreateOAuthCredentialsProvider(const TString& token); + +/* + * Provider that takes IAM ticket from metadata service. + */ +std::shared_ptr<ICredentialsProvider> CreateIAMCredentialsForwarder(TIntrusivePtr<ILogger> logger); + +/* + * Provider that makes IAM ticket from JWT token. + */ +std::shared_ptr<ICredentialsProvider> CreateIAMJwtFileCredentialsForwarder(const TString& jwtKeyFilename, TIntrusivePtr<ILogger> logger, + const TString& endpoint = "iam.api.cloud.yandex.net", const TDuration& refreshPeriod = TDuration::Hours(1), + const TDuration& requestTimeout = TDuration::Seconds(10)); + +std::shared_ptr<ICredentialsProvider> CreateIAMJwtParamsCredentialsForwarder(const TString& jwtParams, TIntrusivePtr<ILogger> logger, + const TString& endpoint = "iam.api.cloud.yandex.net", const TDuration& refreshPeriod = TDuration::Hours(1), + const TDuration& requestTimeout = TDuration::Seconds(10)); +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/iconsumer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/iconsumer.h new file mode 100644 index 0000000000..975b2ccd72 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/iconsumer.h @@ -0,0 +1,39 @@ +#pragma once + +#include "types.h" +#include "responses.h" + +#include <kikimr/yndx/api/protos/persqueue.pb.h> + +#include <library/cpp/threading/future/future.h> + +#include <util/generic/vector.h> + +namespace NPersQueue { + +class IConsumer { +public: + virtual ~IConsumer() = default; + + // Start consumer. + // Consumer can be used after its start will be finished. + virtual NThreading::TFuture<TConsumerCreateResponse> Start(TInstant deadline = TInstant::Max()) noexcept = 0; + + NThreading::TFuture<TConsumerCreateResponse> Start(TDuration timeout) noexcept { + return Start(TInstant::Now() + timeout); + } + + //read result according to settings or Lock/Release requests from server(if UseLockSession is true) + virtual NThreading::TFuture<TConsumerMessage> GetNextMessage() noexcept = 0; + + //commit processed reads + virtual void Commit(const TVector<ui64>& cookies) noexcept = 0; + + //will request status - response will be received in GetNextMessage with EMT_STATUS + virtual void RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept { Y_UNUSED(topic); Y_UNUSED(partition); Y_UNUSED(generation); }; + + //will be signalled in case of errors - you may use this future if you can't use GetNextMessage() //out of memory or any other reason + virtual NThreading::TFuture<TError> IsDead() noexcept = 0; +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/channel.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/channel.cpp new file mode 100644 index 0000000000..c29b1d8dd8 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/channel.cpp @@ -0,0 +1,421 @@ +#include "channel.h" +#include "channel_p.h" +#include "internals.h" +#include "persqueue_p.h" +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> + +#include <ydb/public/sdk/cpp/client/resources/ydb_resources.h> + +#include <util/random/random.h> +#include <util/string/builder.h> + +#include <grpc++/create_channel.h> + +const ui32 MAX_MESSAGE_SIZE = 150 * 1024 * 1024; +const TDuration KEEP_ALIVE_TIME = TDuration::Seconds(90); +const TDuration KEEP_ALIVE_TIMEOUT = TDuration::Seconds(1); +constexpr ui64 MAGIC_COOKIE_VALUE = 123456789; + +namespace NPersQueue { + +void FillMetaHeaders(grpc::ClientContext& context, const TString& database, ICredentialsProvider* credentials) { + context.AddMetadata(NYdb::YDB_AUTH_TICKET_HEADER, GetToken(credentials)); + context.AddMetadata(NYdb::YDB_DATABASE_HEADER, database); +} + +TServerSetting ApplyClusterEndpoint(const TServerSetting& server, const TString& endpoint) { + TServerSetting ret = server; + ret.UseLogbrokerCDS = EClusterDiscoveryUsageMode::DontUse; + TStringBuf endpointWithPort = endpoint; + TStringBuf host, port; + ui16 parsedPort = 0; + if (endpointWithPort.TryRSplit(':', host, port) && TryFromString(port, parsedPort)) { + ret.Address = TString(host); + ret.Port = parsedPort; + } else { + ret.Address = endpoint; + } + return ret; +} + +bool UseCDS(const TServerSetting& server) { + switch (server.UseLogbrokerCDS) { + case EClusterDiscoveryUsageMode::Auto: + return server.Address == "logbroker.yandex.net"sv || server.Address == "logbroker-prestable.yandex.net"sv; + case EClusterDiscoveryUsageMode::Use: + return true; + case EClusterDiscoveryUsageMode::DontUse: + return false; + } +} + +class TChannelImpl::TGetProxyHandler : public IHandler { +public: + TGetProxyHandler(TChannelImplPtr ptr) + : Ptr(ptr) + {} + + void Done() override { + Ptr->OnGetProxyDone(); + } + + void Destroy(const TError& reason) override { + Ptr->Destroy(reason); + } + + TString ToString() override { return "GetProxyHandler"; } + +protected: + TChannelImplPtr Ptr; +}; + +class TChannelOverCdsBaseImpl::TCdsResponseHandler : public IHandler { +public: + TCdsResponseHandler(TChannelOverCdsBaseImplPtr ptr) + : Ptr(ptr) + {} + + void Done() override { + Ptr->OnCdsRequestDone(); + } + + void Destroy(const TError& reason) override { + Ptr->Destroy(reason); + } + + TString ToString() override { return "CDS-ResponseHandler"; } + +protected: + TChannelOverCdsBaseImplPtr Ptr; +}; + +TChannel::TChannel( + const TServerSetting& server, const TCredProviderPtr& credentialsProvider, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger, bool preferLocalProxy +) { + MakeImpl(server, credentialsProvider, pqLib, logger, preferLocalProxy); +} + +TChannel::TChannel( + const TProducerSettings& settings, TPQLibPrivate* pqLib, TIntrusivePtr<ILogger> logger, bool preferLocalProxy +) { + if (UseCDS(settings.Server)) + Impl = new TProducerChannelOverCdsImpl(settings, pqLib, std::move(logger), preferLocalProxy); + else + MakeImpl(settings.Server, settings.CredentialsProvider, pqLib, logger, preferLocalProxy); +} + +void TChannel::MakeImpl( + const TServerSetting& server, const std::shared_ptr<ICredentialsProvider>& credentialsProvider, + TPQLibPrivate* pqLib, TIntrusivePtr<ILogger> logger, bool preferLocalProxy +) { + Impl = MakeIntrusive<TChannelOverDiscoveryImpl>( + server, credentialsProvider, pqLib, std::move(logger), preferLocalProxy + ); +} + +NThreading::TFuture<TChannelInfo> TChannel::GetChannel() { + return Impl->GetChannel(); +} + +TChannel::~TChannel() { + Impl->TryCancel(); + Impl->Wait(); +} + +void TChannel::Start() { + Impl->Start(); +} + +TChannelImpl::TChannelImpl(const TServerSetting& server, const std::shared_ptr<ICredentialsProvider>& credentialsProvider, TPQLibPrivate* pqLib, TIntrusivePtr<ILogger> logger, bool preferLocalProxy) + : Promise(NThreading::NewPromise<TChannelInfo>()) + , Server(server) + , SelectedEndpoint(Server.GetFullAddressString()) + , CredentialsProvider(credentialsProvider) + , PreferLocalProxy(preferLocalProxy) + , Logger(std::move(logger)) + , PQLib(pqLib) + , CQ(pqLib->GetCompletionQueue()) + , ChooseProxyFinished(0) +{} + +TChannelOverDiscoveryImpl::TChannelOverDiscoveryImpl( + const TServerSetting& server, const TCredProviderPtr& credentialsProvider, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger, bool preferLocalProxy +) + : TChannelImpl(server, credentialsProvider, pqLib, logger, preferLocalProxy) +{ + CreationTimeout = PQLib->GetSettings().ChannelCreationTimeout; +} + +TChannelOverCdsBaseImpl::TChannelOverCdsBaseImpl(const TServerSetting& server, const TCredProviderPtr& credentialsProvider, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger, bool preferLocalProxy +) + : TChannelOverDiscoveryImpl(server, credentialsProvider, pqLib, logger, preferLocalProxy) + , CdsRequestFinished(0) +{} + +void TChannelOverDiscoveryImpl::Start() { + Channel = CreateGrpcChannel(SelectedEndpoint); + Stub = Ydb::Discovery::V1::DiscoveryService::NewStub(Channel); + IHandlerPtr handler = new TGetProxyHandler(this); + Context.set_deadline(TInstant::Now() + CreationTimeout); + Ydb::Discovery::ListEndpointsRequest request; + request.set_database(Server.Database); + FillMetaHeaders(Context, Server.Database, CredentialsProvider.get()); + DEBUG_LOG("Send list endpoints request: " << request, "", ""); + Rpc = Stub->AsyncListEndpoints(&Context, request, CQ.get()); + Rpc->Finish(&Response, &Status, new TQueueEvent(std::move(handler))); +} + +void TChannelOverCdsBaseImpl::Start() { + Channel = CreateGrpcChannel(Server.GetFullAddressString()); + Stub = Ydb::PersQueue::V1::ClusterDiscoveryService::NewStub(Channel); + IHandlerPtr handler = new TCdsResponseHandler(this); + CdsContext.set_deadline(TInstant::Now() + PQLib->GetSettings().ChannelCreationTimeout); + Ydb::PersQueue::ClusterDiscovery::DiscoverClustersRequest request = GetCdsRequest(); + FillMetaHeaders(CdsContext, Server.Database, CredentialsProvider.get()); + DEBUG_LOG("Send cds request: " << request, "", ""); + CreationStartTime = TInstant::Now(); + Rpc = Stub->AsyncDiscoverClusters(&CdsContext, request, CQ.get()); + Rpc->Finish(&Response, &Status, new TQueueEvent(std::move(handler))); +} + +NThreading::TFuture<TChannelInfo> TChannelImpl::GetChannel() { + return Promise.GetFuture(); +} + +bool TChannelImpl::TryCancel() { + if (AtomicSwap(&ChooseProxyFinished, 1) == 0) { + Y_ASSERT(!Promise.HasValue()); + Promise.SetValue(TChannelInfo{nullptr, 0}); + Context.TryCancel(); + return true; + } + return false; +} + +TChannelImpl::~TChannelImpl() = default; + +std::shared_ptr<grpc::Channel> TChannelImpl::CreateGrpcChannel(const TString& address) { + grpc::ChannelArguments args; + args.SetInt(GRPC_ARG_KEEPALIVE_TIME_MS, KEEP_ALIVE_TIME.MilliSeconds()); + args.SetInt(GRPC_ARG_KEEPALIVE_TIMEOUT_MS, KEEP_ALIVE_TIMEOUT.MilliSeconds()); + args.SetInt(GRPC_ARG_KEEPALIVE_PERMIT_WITHOUT_CALLS, 1); + args.SetInt(GRPC_ARG_HTTP2_MAX_PINGS_WITHOUT_DATA, 0); + args.SetInt(GRPC_ARG_HTTP2_MIN_SENT_PING_INTERVAL_WITHOUT_DATA_MS, KEEP_ALIVE_TIME.MilliSeconds()); + + args.SetInt(GRPC_ARG_MAX_RECEIVE_MESSAGE_LENGTH, MAX_MESSAGE_SIZE); + args.SetInt(GRPC_ARG_MAX_SEND_MESSAGE_LENGTH, MAX_MESSAGE_SIZE); + + //TODO: fill there other settings from PQLib + DEBUG_LOG("Creating grpc channel to \"" << address << "\"", "", ""); + if (Server.UseSecureConnection) + return grpc::CreateCustomChannel(address, + grpc::SslCredentials(grpc::SslCredentialsOptions{Server.CaCert, "", ""}), + args); + else + return grpc::CreateCustomChannel(address, grpc::InsecureChannelCredentials(), args); +} + +TString TChannelImpl::GetProxyAddress(const TString& proxyName, ui32 port) { + TStringBuilder address; + if (proxyName.find(':') != TString::npos && proxyName.find('[') == TString::npos) { //fix for ipv6 adresses + address << "[" << proxyName << "]"; + } else { + address << proxyName; + } + if (port) + address << ":" << port; + else + address << ":" << Server.Port; + return std::move(address); +} + +void TChannelImpl::OnGetProxyDone() { + if (AtomicSwap(&ChooseProxyFinished, 1) == 0) { + const bool ok = Status.ok(); + Y_ASSERT(!Promise.HasValue()); + if (ok) { + ProcessGetProxyResponse(); + } else { + if (Status.error_code() == grpc::CANCELLED) { + INFO_LOG("Grpc request is canceled by user", "", ""); + } else { + const auto& msg = Status.error_message(); + TString res(msg.data(), msg.length()); + ERR_LOG("Grpc error " << static_cast<int>(Status.error_code()) << ": \"" << res << "\"", "", ""); + } + StartFailed(); + } + } + SetDone(); +} + +void TChannelImpl::StartFailed() { + AtomicSet(ChooseProxyFinished, 1); + Promise.SetValue(TChannelInfo{nullptr, 0}); +} + +void TChannelOverDiscoveryImpl::ProcessGetProxyResponse() { + Ydb::Discovery::ListEndpointsResult result; + auto& reply = Response.operation().result(); + reply.UnpackTo(&result); + + auto has_service = [](const Ydb::Discovery::EndpointInfo& endpoint) { + for (auto& service: endpoint.service()) { + if (service == "pq") + return true; + } + return false; + }; + TVector<std::pair<TString, ui32>> validEps; + for (const auto& ep : result.endpoints()) { + if (has_service(ep) && (!Server.UseSecureConnection || ep.ssl())) { + validEps.emplace_back(ep.address(), ep.port()); + } + } + if (validEps.empty()) { + ERR_LOG("No valid endpoints to connect!", "", ""); + StartFailed(); + } else { + auto& selected = validEps[RandomNumber<ui32>(validEps.size())]; + DEBUG_LOG("Selected endpoint {\"" << selected.first << "\", " << selected.second << "}", "", ""); + auto channel = CreateGrpcChannel(GetProxyAddress(selected.first, selected.second)); + Promise.SetValue(TChannelInfo{std::move(channel), MAGIC_COOKIE_VALUE}); + } +} + +void TChannelOverCdsBaseImpl::OnCdsRequestDone() { + DEBUG_LOG("ON CDS request done: " << Response, "", ""); + + if (AtomicSwap(&CdsRequestFinished, 1) == 0) { + const bool ok = Status.ok(); + if (ok) { + if (ProcessCdsResponse()) { + CreationTimeout = CreationTimeout - (TInstant::Now() - CreationStartTime); + TChannelOverDiscoveryImpl::Start(); + // Don't set AllDone event, because we are about to continue node discovery + // via TChannelOverDiscoveryImpl part of implementation. + return; + } else { + StartFailed(); + } + } else { + if (Status.error_code() == grpc::CANCELLED) { + INFO_LOG("Grpc request is canceled by user", "", ""); + } else { + const auto& msg = Status.error_message(); + TString res(msg.data(), msg.length()); + ERR_LOG("Grpc error " << static_cast<int>(Status.error_code()) << ": \"" << res << "\"", "", ""); + } + StartFailed(); + } + } + SetDone(); +} + +bool TChannelOverCdsBaseImpl::TryCancel() { + if (AtomicSwap(&CdsRequestFinished, 1) == 0) { + Y_ASSERT(!Promise.HasValue()); + Promise.SetValue(TChannelInfo{nullptr, 0}); + CdsContext.TryCancel(); + return true; + } else { + TChannelOverDiscoveryImpl::TryCancel(); + } + return false; +} + +TProducerChannelOverCdsImpl::TProducerChannelOverCdsImpl(const TProducerSettings& settings, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger, bool preferLocalProxy) + : TChannelOverCdsBaseImpl(settings.Server, settings.CredentialsProvider, pqLib, logger, preferLocalProxy) + , Topic(settings.Topic) + , SourceId(settings.SourceId) + , PreferredCluster(settings.PreferredCluster) +{ +} + +Ydb::PersQueue::ClusterDiscovery::DiscoverClustersRequest TProducerChannelOverCdsImpl::GetCdsRequest() const { + Ydb::PersQueue::ClusterDiscovery::DiscoverClustersRequest request; + auto* params = request.add_write_sessions(); + params->set_topic(Topic); + params->set_source_id(SourceId); + params->set_preferred_cluster_name(PreferredCluster); // Can be empty. + return request; +} + +bool TProducerChannelOverCdsImpl::ProcessCdsResponse() { + Ydb::PersQueue::ClusterDiscovery::DiscoverClustersResult result; + auto& reply = Response.operation().result(); + reply.UnpackTo(&result); + DEBUG_LOG("Process CDS result: " << result, "", ""); + for (auto& session: result.write_sessions_clusters()) { + for (auto& cluster: session.clusters()) { + if (cluster.available()) { + const TServerSetting selectedServer = ApplyClusterEndpoint(Server, cluster.endpoint()); + SelectedEndpoint = GetProxyAddress(selectedServer.Address, selectedServer.Port); + return true; + } + } + } + ERR_LOG("Could not find valid cluster in CDS response", "", ""); + return false; +} + +TConsumerChannelOverCdsImpl::TConsumerChannelOverCdsImpl(const TConsumerSettings& settings, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger) + : TChannelOverCdsBaseImpl(settings.Server, settings.CredentialsProvider, pqLib, logger, true) + , Result(std::make_shared<TResult>()) + , Topics(settings.Topics) +{ +} + +Ydb::PersQueue::ClusterDiscovery::DiscoverClustersRequest TConsumerChannelOverCdsImpl::GetCdsRequest() const { + Ydb::PersQueue::ClusterDiscovery::DiscoverClustersRequest request; + for (const TString& topic : Topics) { + auto* params = request.add_read_sessions(); + params->set_topic(topic); + params->mutable_all_original(); // set all_original + } + return request; +} + +bool TConsumerChannelOverCdsImpl::ProcessCdsResponse() { + auto& reply = Response.operation().result(); + reply.UnpackTo(&Result->second); + Result->first = Status; + DEBUG_LOG("Got CDS result: " << Result->second, "", ""); + return false; // Signals future +} + +void TConsumerChannelOverCdsImpl::StartFailed() { + Result->first = Status; + TChannelOverCdsBaseImpl::StartFailed(); +} + +void TChannelImpl::SetDone() { + AllDoneEvent.Signal(); +} + +void TChannelImpl::Destroy(const TError& error) { + if (AtomicSwap(&ChooseProxyFinished, 1) == 0) { + const auto& msg = Status.error_message(); + TString res(msg.data(), msg.length()); + + ERR_LOG("Got proxy error response " << error << " grpc error " << res, "", ""); + //TODO: add here status from Status to response, log status + + Y_ASSERT(!Promise.HasValue()); + Promise.SetValue(TChannelInfo{nullptr, 0}); + } + + SetDone(); +} + +void TChannelImpl::Wait() { + AllDoneEvent.Wait(); +} + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/channel.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/channel.h new file mode 100644 index 0000000000..630b097675 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/channel.h @@ -0,0 +1,50 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> + +#include <library/cpp/threading/future/future.h> + +#include <deque> + +namespace NPersQueue { + +struct TChannelInfo { + std::shared_ptr<grpc::Channel> Channel; + ui64 ProxyCookie; +}; + +class TChannel; +using TChannelPtr = TIntrusivePtr<TChannel>; + +class TPQLibPrivate; + +struct TChannelHolder { + TChannelPtr ChannelPtr; + NThreading::TFuture<TChannelInfo> ChannelInfo; +}; + +class TChannelImpl; +using TChannelImplPtr = TIntrusivePtr<TChannelImpl>; + +class TChannel: public TAtomicRefCount<TChannel> { +public: + friend class TPQLibPrivate; + + NThreading::TFuture<TChannelInfo> GetChannel(); + + ~TChannel(); + + void Start(); + +private: + TChannel(const TServerSetting& server, const std::shared_ptr<ICredentialsProvider>& credentialsProvider, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger = nullptr, bool preferLocalProxy = false); + TChannel(const TProducerSettings& settings, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger = nullptr, bool preferLocalProxy = false); + void MakeImpl(const TServerSetting& server, const TCredProviderPtr&, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger = nullptr, bool preferLocalProxy = false); + TChannelImplPtr Impl; +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/channel_p.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/channel_p.h new file mode 100644 index 0000000000..10f18abb04 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/channel_p.h @@ -0,0 +1,156 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include "channel.h" +#include "scheduler.h" +#include "internals.h" +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> +#include <library/cpp/threading/future/future.h> +#include <ydb/public/api/protos/ydb_discovery.pb.h> +#include <ydb/public/api/grpc/ydb_discovery_v1.grpc.pb.h> +#include <ydb/public/api/grpc/draft/ydb_persqueue_v1.grpc.pb.h> + +#include <deque> + +namespace NPersQueue { + +class TPQLibPrivate; + +// Apply endpoint that was taken from CDS. +TServerSetting ApplyClusterEndpoint(const TServerSetting& server, const TString& endpoint); + +class TChannelImpl : public TAtomicRefCount<TChannelImpl> { +public: + NThreading::TFuture<TChannelInfo> GetChannel(); + + virtual ~TChannelImpl(); + + virtual bool TryCancel(); + + TChannelImpl(const TServerSetting& server, const TCredProviderPtr& credentialsProvider, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger, bool preferLocalProxy); + + virtual void Start() = 0; + + void Wait(); + +protected: + friend class TPQLibPrivate; + class TGetProxyHandler; + + void Destroy(const TError& error); + virtual void StartFailed(); + void SetDone(); + void OnGetProxyDone(); + virtual void ProcessGetProxyResponse() = 0; + std::shared_ptr<grpc::Channel> CreateGrpcChannel(const TString& address); + TString GetProxyAddress(const TString& proxyName, ui32 port = 0); + TString GetToken(); + + NThreading::TPromise<TChannelInfo> Promise; + TServerSetting Server; + TString SelectedEndpoint; + std::shared_ptr<ICredentialsProvider> CredentialsProvider; + bool PreferLocalProxy; + + TIntrusivePtr<ILogger> Logger; + TPQLibPrivate* PQLib; + +protected: + std::shared_ptr<grpc::CompletionQueue> CQ; + std::shared_ptr<grpc::Channel> Channel; + grpc::ClientContext Context; + grpc::Status Status; + + TAutoEvent AllDoneEvent; + + TAtomic ChooseProxyFinished; +}; + +class TChannelOverDiscoveryImpl : public TChannelImpl { +public: + TChannelOverDiscoveryImpl(const TServerSetting& server, const TCredProviderPtr&, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger, bool preferLocalProxy); + void Start() override; + +protected: + void ProcessGetProxyResponse() override; + TDuration CreationTimeout; + +private: + std::unique_ptr<Ydb::Discovery::V1::DiscoveryService::Stub> Stub; + std::unique_ptr<grpc::ClientAsyncResponseReader<Ydb::Discovery::ListEndpointsResponse>> Rpc; + Ydb::Discovery::ListEndpointsResponse Response; +}; + + +class TChannelOverCdsBaseImpl : public TChannelOverDiscoveryImpl { +public: + TChannelOverCdsBaseImpl(const TServerSetting& server, const TCredProviderPtr&, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger, bool preferLocalProxy); + void Start() override; + void OnCdsRequestDone(); + + bool TryCancel() override; + +private: + virtual Ydb::PersQueue::ClusterDiscovery::DiscoverClustersRequest GetCdsRequest() const = 0; + virtual bool ProcessCdsResponse() = 0; + +protected: + class TCdsResponseHandler; + +private: + std::unique_ptr<Ydb::PersQueue::V1::ClusterDiscoveryService::Stub> Stub; + std::unique_ptr<grpc::ClientAsyncResponseReader<Ydb::PersQueue::ClusterDiscovery::DiscoverClustersResponse>> Rpc; + +protected: + Ydb::PersQueue::ClusterDiscovery::DiscoverClustersResponse Response; + +private: + grpc::ClientContext CdsContext; + + TInstant CreationStartTime = TInstant::Zero(); + TAtomic CdsRequestFinished; +}; +using TChannelOverCdsBaseImplPtr = TIntrusivePtr<TChannelOverCdsBaseImpl>; + +class TProducerChannelOverCdsImpl : public TChannelOverCdsBaseImpl { +public: + TProducerChannelOverCdsImpl(const TProducerSettings& settings, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger, bool preferLocalProxy); + +private: + Ydb::PersQueue::ClusterDiscovery::DiscoverClustersRequest GetCdsRequest() const override; + bool ProcessCdsResponse() override; + +private: + TString Topic; + TString SourceId; + TString PreferredCluster; +}; + +class TConsumerChannelOverCdsImpl : public TChannelOverCdsBaseImpl { +public: + using TResult = std::pair<grpc::Status, Ydb::PersQueue::ClusterDiscovery::DiscoverClustersResult>; + using TResultPtr = std::shared_ptr<TResult>; + +public: + TConsumerChannelOverCdsImpl(const TConsumerSettings& settings, TPQLibPrivate* pqLib, + TIntrusivePtr<ILogger> logger); + + TResultPtr GetResultPtr() const { + return Result; + } + +private: + Ydb::PersQueue::ClusterDiscovery::DiscoverClustersRequest GetCdsRequest() const override; + bool ProcessCdsResponse() override; + void StartFailed() override; + +private: + TResultPtr Result; + TVector<TString> Topics; +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compat_producer.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compat_producer.cpp new file mode 100644 index 0000000000..14c2871a54 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compat_producer.cpp @@ -0,0 +1,256 @@ + +#include "compat_producer.h" +#include "persqueue_p.h" + +namespace NPersQueue { + + using namespace NYdb::NPersQueue; + + + static TWriteSessionSettings ConvertToWriteSessionSettings(const TProducerSettings& ps) { + TWriteSessionSettings settings; + settings.Path(ps.Topic) + .MessageGroupId(ps.SourceId) + .RetryPolicy(ps.ReconnectOnFailure ? NYdb::NPersQueue::IRetryPolicy::GetExponentialBackoffPolicy(ps.ReconnectionDelay, ps.ReconnectionDelay, ps.MaxReconnectionDelay, + ps.MaxAttempts, TDuration::Max(), 2.0, [](NYdb::EStatus){ return ERetryErrorClass::ShortRetry; }) + : NYdb::NPersQueue::IRetryPolicy::GetNoRetryPolicy()) + .Codec(NYdb::NPersQueue::ECodec(ps.Codec + 1)) + .CompressionLevel(ps.Quality) + .ValidateSeqNo(false); + for (auto& attr : ps.ExtraAttrs) { + settings.AppendSessionMeta(attr.first, attr.second); + } + + if (ps.PreferredCluster) { + settings.PreferredCluster(ps.PreferredCluster); + } + if (ps.DisableCDS || ps.Server.UseLogbrokerCDS == EClusterDiscoveryUsageMode::DontUse) { + settings.ClusterDiscoveryMode(NYdb::NPersQueue::EClusterDiscoveryMode::Off); + } else if (ps.Server.UseLogbrokerCDS == EClusterDiscoveryUsageMode::Use) { + settings.ClusterDiscoveryMode(NYdb::NPersQueue::EClusterDiscoveryMode::On); + } + + return settings; + } + + TYdbSdkCompatibilityProducer::~TYdbSdkCompatibilityProducer() { + WriteSession->Close(TDuration::Seconds(0)); + + NotifyClient(NErrorCode::OK, "producer object destroyed by client"); + } + + TYdbSdkCompatibilityProducer::TYdbSdkCompatibilityProducer(const TProducerSettings& settings, TPersQueueClient& persQueueClient, + std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib) + : IProducerImpl(destroyEventRef, pqLib) + { + IsDeadPromise = ::NThreading::NewPromise<TError>(); + TWriteSessionSettings sessionSettings = ConvertToWriteSessionSettings(settings); + Closed = false; + WriteSession = persQueueClient.CreateWriteSession(sessionSettings); + } + + void TYdbSdkCompatibilityProducer::SubscribeToNextEvent() { + NextEvent = WriteSession->WaitEvent(); + std::weak_ptr<TYdbSdkCompatibilityProducer> self = shared_from_this(); + NextEvent.Subscribe([self](const auto& future) { + auto selfShared = self.lock(); + Y_UNUSED(future); + if (selfShared) { + selfShared->DoProcessNextEvent(); + } + }); + } + + void TYdbSdkCompatibilityProducer::Init() noexcept { + } + + NThreading::TFuture<TProducerCreateResponse> TYdbSdkCompatibilityProducer::Start(TInstant deadline) noexcept { + Y_UNUSED(deadline); // TODO: start timeout? + SubscribeToNextEvent(); + return WriteSession->GetInitSeqNo().Apply([](const auto& future) { + TWriteResponse response; + if (future.HasException()) { + response.MutableError()->SetDescription("session closed"); + response.MutableError()->SetCode(NErrorCode::ERROR); + } else { + response.MutableInit()->SetMaxSeqNo(future.GetValue()); + } + TProducerCreateResponse res(std::move(response)); + + return NThreading::MakeFuture<TProducerCreateResponse>(res); + }); + } + + void TYdbSdkCompatibilityProducer::WriteImpl(NThreading::TPromise<TProducerCommitResponse>& promise, TMaybe<TProducerSeqNo> seqNo, ::NPersQueue::TData data) noexcept { + TContinuationToken *contToken = nullptr; + bool write = false; + with_lock(Lock) { + if (!Closed) { + contToken = ContToken.Get(); + if(contToken == nullptr || !ToWrite.empty()) { + ToWrite.emplace(promise, seqNo, data); + return; + } else { // contToken here and ToWrite is empty - can write right now + ToAck.emplace(promise, seqNo, data); + write = true; + } + } + } + if (write) { + if (data.IsEncoded()) { + WriteSession->WriteEncoded(std::move(*contToken), data.GetEncodedData(), NYdb::NPersQueue::ECodec(data.GetCodecType() + 1), data.GetOriginalSize(), seqNo); // can trigger NextEvent.Subscribe() in the same thread + } else { + WriteSession->Write(std::move(*contToken), data.GetSourceData(), seqNo); // can trigger NextEvent.Subscribe() in the same thread + } + } else { + TWriteResponse wr; + wr.MutableError()->SetCode(NErrorCode::ERROR); + wr.MutableError()->SetDescription("session closed"); + TProducerCommitResponse resp(seqNo ? *seqNo : 0, data, std::move(wr)); + promise.SetValue(resp); + } + } + + void TYdbSdkCompatibilityProducer::Write(NThreading::TPromise<TProducerCommitResponse>& promise, ::NPersQueue::TData data) noexcept { + return WriteImpl(promise, {}, std::move(data)); // similar to TCompressingProducer + } + + void TYdbSdkCompatibilityProducer::Write(NThreading::TPromise<TProducerCommitResponse>& promise, TProducerSeqNo seqNo, ::NPersQueue::TData data) noexcept { + return WriteImpl(promise, seqNo, std::move(data)); // similar to TCompressingProducer + } + + void TYdbSdkCompatibilityProducer::Cancel() { + WriteSession->Close(TDuration::Seconds(0)); + } + + NErrorCode::EErrorCode GetErrorCode(const NYdb::EStatus status) { + switch(status) { + case NYdb::EStatus::SUCCESS: + return NErrorCode::OK; + case NYdb::EStatus::UNAVAILABLE: + return NErrorCode::INITIALIZING; + case NYdb::EStatus::OVERLOADED: + return NErrorCode::OVERLOAD; + case NYdb::EStatus::BAD_REQUEST: + return NErrorCode::BAD_REQUEST; + case NYdb::EStatus::NOT_FOUND: + case NYdb::EStatus::SCHEME_ERROR: + return NErrorCode::UNKNOWN_TOPIC; + case NYdb::EStatus::UNSUPPORTED: + return NErrorCode::BAD_REQUEST; + case NYdb::EStatus::UNAUTHORIZED: + return NErrorCode::ACCESS_DENIED; + default: + return NErrorCode::ERROR; + } + return NErrorCode::ERROR; + } + + struct TToNotify { + NThreading::TPromise<TProducerCommitResponse> Promise; + TProducerCommitResponse Resp; + + TToNotify(NThreading::TPromise<TProducerCommitResponse>& promise, TProducerCommitResponse&& resp) + : Promise(promise) + , Resp(resp) + {} + }; + + + void TYdbSdkCompatibilityProducer::NotifyClient(NErrorCode::EErrorCode code, const TString& reason) { + + TError err; + err.SetCode(code); + err.SetDescription(reason); + std::queue<TToNotify> toNotify; + with_lock(Lock) { + if (Closed) return; + Closed = true; + while(!ToAck.empty()) { + TMsgData& item = ToAck.front(); + TWriteResponse wr; + wr.MutableError()->CopyFrom(err); + TProducerCommitResponse resp(item.SeqNo ? *item.SeqNo : 0, item.Data, std::move(wr)); + toNotify.emplace(item.Promise, std::move(resp)); + ToAck.pop(); + } + while(!ToWrite.empty()) { + TMsgData& item = ToWrite.front(); + TWriteResponse wr; + wr.MutableError()->CopyFrom(err); + TProducerCommitResponse resp(item.SeqNo ? *item.SeqNo : 0, item.Data, std::move(wr)); + toNotify.emplace(item.Promise, std::move(resp)); + ToWrite.pop(); + } + } + while(!toNotify.empty()) { + toNotify.front().Promise.SetValue(toNotify.front().Resp); + toNotify.pop(); + } + IsDeadPromise.SetValue(err); + DestroyPQLibRef(); + } + + void TYdbSdkCompatibilityProducer::DoProcessNextEvent() { + + std::queue<TToNotify> toNotify; + + TVector<TWriteSessionEvent::TEvent> events = WriteSession->GetEvents(false); + //Y_VERIFY(!events.empty()); + for (auto& event : events) { + if(std::holds_alternative<TSessionClosedEvent>(event)) { + NotifyClient(GetErrorCode(std::get<TSessionClosedEvent>(event).GetStatus()), std::get<TSessionClosedEvent>(event).DebugString() ); + return; + } else if(std::holds_alternative<TWriteSessionEvent::TAcksEvent>(event)) { + + // get seqNos, signal their promises + TWriteSessionEvent::TAcksEvent& acksEvent = std::get<TWriteSessionEvent::TAcksEvent>(event); + std::queue<TToNotify> toNotify; + for(TWriteSessionEvent::TWriteAck& ackEvent : acksEvent.Acks) { + TWriteResponse writeResp; + auto ackMsg = writeResp.MutableAck(); + ackMsg->SetSeqNo(ackEvent.SeqNo); + ackMsg->SetAlreadyWritten(ackEvent.State == TWriteSessionEvent::TWriteAck::EEventState::EES_ALREADY_WRITTEN); + ackMsg->SetOffset(ackEvent.Details ? ackEvent.Details->Offset : 0); + with_lock(Lock) { + Y_ASSERT(!ToAck.empty()); + TProducerCommitResponse resp(ackEvent.SeqNo, ToAck.front().Data, std::move(writeResp)); + toNotify.emplace(ToAck.front().Promise, std::move(resp)); + ToAck.pop(); + } + } + while(!toNotify.empty()) { + toNotify.front().Promise.SetValue(toNotify.front().Resp); + toNotify.pop(); + } + } else if(std::holds_alternative<TWriteSessionEvent::TReadyToAcceptEvent>(event)) { + TWriteSessionEvent::TReadyToAcceptEvent& readyEvent = std::get<TWriteSessionEvent::TReadyToAcceptEvent>(event); + TContinuationToken contToken = std::move(readyEvent.ContinuationToken); // IF there is only one only movable contToken THEN this is thread safe + TString strData; + TMaybe<TProducerSeqNo> seqNo; + bool write = false; + with_lock(Lock) { + if(ToWrite.empty()) { + ContToken = std::move(contToken); + } else { + strData = ToWrite.front().Data.GetSourceData(); + seqNo = ToWrite.front().SeqNo; + ToAck.push(ToWrite.front()); + ToWrite.pop(); + write = true; + } + } + if (write) { + WriteSession->Write(std::move(contToken), strData, seqNo); + } + } + } + + SubscribeToNextEvent(); + } + + NThreading::TFuture<TError> TYdbSdkCompatibilityProducer::IsDead() noexcept { + return IsDeadPromise.GetFuture(); + } + +} // namespace NYdb::NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compat_producer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compat_producer.h new file mode 100644 index 0000000000..e5287d4694 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compat_producer.h @@ -0,0 +1,66 @@ +#pragma once + +#include "iproducer_p.h" + +#include <kikimr/public/sdk/cpp/client/ydb_persqueue/persqueue.h> + +#include <memory> +#include <queue> + +namespace NPersQueue { + + +/** + * Compatibility producer: + * wraps old "interface" around the new one + */ + +class TYdbSdkCompatibilityProducer: public IProducerImpl, public std::enable_shared_from_this<TYdbSdkCompatibilityProducer> { + + struct TMsgData { + NThreading::TPromise<TProducerCommitResponse> Promise; + TMaybe<TProducerSeqNo> SeqNo; + TData Data; + + TMsgData() {}; + TMsgData(NThreading::TPromise<TProducerCommitResponse> promise, TMaybe<TProducerSeqNo> seqNo, TData data) + : Promise(promise) + , SeqNo(seqNo) + , Data(data) {} + }; + + std::shared_ptr<NYdb::NPersQueue::IWriteSession> WriteSession; + NThreading::TPromise<TError> IsDeadPromise; + NThreading::TFuture<void> NextEvent; + + TMaybe<NYdb::NPersQueue::TContinuationToken> ContToken; + std::queue<TMsgData> ToWrite; + std::queue<TMsgData> ToAck; + TSpinLock Lock; + bool Closed; + + void DoProcessNextEvent(); + void SubscribeToNextEvent(); + + void WriteImpl(NThreading::TPromise<TProducerCommitResponse>& promise, TMaybe<TProducerSeqNo> seqNo, TData data) noexcept; + + void NotifyClient(NErrorCode::EErrorCode code, const TString& reason); + +public: + using IProducerImpl::Write; + NThreading::TFuture<TProducerCreateResponse> Start(TInstant deadline = TInstant::Max()) noexcept override; + NThreading::TFuture<TError> IsDead() noexcept override; + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TData data) noexcept override; + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TProducerSeqNo seqNo, TData data) noexcept override; + void Cancel() override; + void Init() noexcept override; + + ~TYdbSdkCompatibilityProducer(); + + TYdbSdkCompatibilityProducer(const TProducerSettings& settings, NYdb::NPersQueue::TPersQueueClient& persQueueClient, + std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib); + +}; + + +} // namespace NYdb::NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compressing_producer.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compressing_producer.cpp new file mode 100644 index 0000000000..ecb6f70944 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compressing_producer.cpp @@ -0,0 +1,155 @@ +#include "compressing_producer.h" +#include "persqueue_p.h" + +namespace NPersQueue { + +TCompressingProducer::TCompressingProducer(std::shared_ptr<IProducerImpl> subproducer, ECodec defaultCodec, int quality, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger) + : IProducerImpl(std::move(destroyEventRef), std::move(pqLib)) + , Logger(std::move(logger)) + , Subproducer(std::move(subproducer)) + , DefaultCodec(defaultCodec) + , Quality(quality) +{ + if (defaultCodec == ECodec::DEFAULT) { + ythrow yexception() << "Producer codec can't be ECodec::DEFAULT"; + } +} + +TCompressingProducer::~TCompressingProducer() { + DestroyQueue("Destructor called"); +} + +NThreading::TFuture<TProducerCreateResponse> TCompressingProducer::Start(TInstant deadline) noexcept { + return Subproducer->Start(deadline); +} + +NThreading::TFuture<TData> TCompressingProducer::Enqueue(TData data) { + NThreading::TPromise<TData> promise = NThreading::NewPromise<TData>(); + std::weak_ptr<TCompressingProducer> self = shared_from_this(); + const void* queueTag = this; + const ECodec defaultCodec = DefaultCodec; + const int quality = Quality; + auto compress = [promise, self, data, queueTag, defaultCodec, quality, pqLib = PQLib.Get()]() mutable { + Compress(promise, std::move(data), defaultCodec, quality, queueTag, self, pqLib); + }; + Y_VERIFY(PQLib->GetCompressionPool().AddFunc(compress)); + return promise.GetFuture(); +} + +void TCompressingProducer::Compress(NThreading::TPromise<TData>& promise, TData&& data, ECodec defaultCodec, int quality, const void* queueTag, std::weak_ptr<TCompressingProducer> self, TPQLibPrivate* pqLib) { + if (!self.lock()) { + return; + } + promise.SetValue(TData::Encode(std::move(data), defaultCodec, quality)); + SignalProcessQueue(queueTag, self, pqLib); +} + +void TCompressingProducer::SignalProcessQueue(const void* queueTag, std::weak_ptr<TCompressingProducer> self, TPQLibPrivate* pqLib) { + auto processQueue = [self] { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->ProcessQueue(); + } + }; + + //can't use this->Logger here - weak_ptr::lock -> ~TCompressingProducer in compression thread -> + //race with cancel() in main thread pool + Y_UNUSED(pqLib->GetQueuePool().GetQueue(queueTag).AddFunc(processQueue)); +} + +void TCompressingProducer::ProcessQueue() { + if (DestroyedFlag) { + return; + } + + size_t removeCount = 0; + bool prevSignalled = true; // flag to guarantee correct order of futures signalling. + for (TWriteRequestInfo& request : Queue) { + if (!request.EncodedData.HasValue()) { + // stage 1 + break; + } + if (!request.Future.Initialized()) { + // stage 2 + if (request.SeqNo) { + request.Future = Subproducer->Write(request.SeqNo, request.EncodedData.GetValue()); + } else { + request.Future = Subproducer->Write(request.EncodedData.GetValue()); + } + std::weak_ptr<TCompressingProducer> self = shared_from_this(); + const void* queueTag = this; + auto signalProcess = [self, queueTag, pqLib = PQLib.Get()](const auto&) { + SignalProcessQueue(queueTag, self, pqLib); + }; + request.Future.Subscribe(signalProcess); + } else { + // stage 3 + if (prevSignalled && request.Future.HasValue()) { + request.Promise.SetValue(request.Future.GetValue()); + ++removeCount; + } else { + prevSignalled = false; + } + } + } + while (removeCount--) { + Queue.pop_front(); + } +} + +void TCompressingProducer::DestroyQueue(const TString& reason) { + if (DestroyedFlag) { + return; + } + + DestroyedFlag = true; + for (TWriteRequestInfo& request : Queue) { + if (request.Future.Initialized() && request.Future.HasValue()) { + request.Promise.SetValue(request.Future.GetValue()); + } else { + TWriteResponse resp; + TError& err = *resp.MutableError(); + err.SetCode(NErrorCode::ERROR); + err.SetDescription(reason); + request.Promise.SetValue(TProducerCommitResponse(request.SeqNo, std::move(request.Data), std::move(resp))); + } + } + Queue.clear(); + + DestroyPQLibRef(); +} + +void TCompressingProducer::Write(NThreading::TPromise<TProducerCommitResponse>& promise, TProducerSeqNo seqNo, TData data) noexcept { + if (DestroyedFlag) { + return; + } + + Queue.emplace_back(seqNo, data, promise); + if (data.IsEncoded()) { + Queue.back().EncodedData = NThreading::MakeFuture<TData>(data); + ProcessQueue(); + } else if (DefaultCodec == ECodec::RAW) { + Queue.back().EncodedData = NThreading::MakeFuture<TData>(TData::MakeRawIfNotEncoded(data)); + ProcessQueue(); + } else { + Queue.back().EncodedData = Enqueue(std::move(data)); + } +} + +void TCompressingProducer::Write(NThreading::TPromise<TProducerCommitResponse>& promise, TData data) noexcept { + TCompressingProducer::Write(promise, 0, std::move(data)); +} + +NThreading::TFuture<TError> TCompressingProducer::IsDead() noexcept { + return Subproducer->IsDead(); +} + +NThreading::TFuture<void> TCompressingProducer::Destroyed() noexcept { + return WaitExceptionOrAll(DestroyedPromise.GetFuture(), Subproducer->Destroyed()); +} + +void TCompressingProducer::Cancel() { + DestroyQueue(GetCancelReason()); +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compressing_producer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compressing_producer.h new file mode 100644 index 0000000000..6ed9e89bed --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compressing_producer.h @@ -0,0 +1,64 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iproducer.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> +#include "internals.h" +#include "iproducer_p.h" + +#include <deque> +#include <memory> + +namespace NPersQueue { + +class TPQLibPrivate; + +class TCompressingProducer: public IProducerImpl, public std::enable_shared_from_this<TCompressingProducer> { +public: + NThreading::TFuture<TProducerCreateResponse> Start(TInstant deadline) noexcept override; + using IProducerImpl::Write; + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TProducerSeqNo seqNo, TData data) noexcept override; + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TData data) noexcept override; + NThreading::TFuture<TError> IsDead() noexcept override; + NThreading::TFuture<void> Destroyed() noexcept override; + + ~TCompressingProducer(); + + TCompressingProducer(std::shared_ptr<IProducerImpl> subproducer, ECodec defaultCodec, int quality, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger); + + void Cancel() override; + +private: + struct TWriteRequestInfo { + TWriteRequestInfo(TProducerSeqNo seqNo, TData data, const NThreading::TPromise<TProducerCommitResponse>& promise) + : SeqNo(seqNo) + , Data(std::move(data)) + , Promise(promise) + { + } + + TProducerSeqNo SeqNo; + TData Data; + NThreading::TFuture<TData> EncodedData; + NThreading::TPromise<TProducerCommitResponse> Promise; + NThreading::TFuture<TProducerCommitResponse> Future; + }; + + void ProcessQueue(); + NThreading::TFuture<TData> Enqueue(TData data); + // these functions should be static and not lock self in compression thread pool for proper futures ordering + static void Compress(NThreading::TPromise<TData>& promise, TData&& data, ECodec defaultCodec, int quality, + const void* queueTag, std::weak_ptr<TCompressingProducer> self, TPQLibPrivate* pqLib); + static void SignalProcessQueue(const void* queueTag, std::weak_ptr<TCompressingProducer> self, TPQLibPrivate* pqLib); + void DestroyQueue(const TString& reason); + +protected: + TIntrusivePtr<ILogger> Logger; + std::shared_ptr<IProducerImpl> Subproducer; + ECodec DefaultCodec; + int Quality; + std::deque<TWriteRequestInfo> Queue; + bool DestroyedFlag = false; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compressing_producer_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compressing_producer_ut.cpp new file mode 100644 index 0000000000..4d9b30f69a --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/compressing_producer_ut.cpp @@ -0,0 +1,243 @@ +#include "compressing_producer.h" +#include "local_caller.h" +#include "persqueue_p.h" + +#include <library/cpp/testing/gmock_in_unittest/gmock.h> +#include <library/cpp/testing/unittest/registar.h> + +#include <util/system/thread.h> + +using namespace testing; + +namespace NPersQueue { + +class TMockProducer: public IProducerImpl { +public: + TMockProducer() + : IProducerImpl(nullptr, nullptr) + { + } + + NThreading::TFuture<TProducerCreateResponse> Start(TInstant) noexcept override { + return NThreading::MakeFuture<TProducerCreateResponse>(MockStart()); + } + + NThreading::TFuture<TProducerCommitResponse> Write(TProducerSeqNo seqNo, TData data) noexcept override { + return NThreading::MakeFuture<TProducerCommitResponse>(MockWriteWithSeqNo(seqNo, data)); + } + + NThreading::TFuture<TProducerCommitResponse> Write(TData data) noexcept override { + return NThreading::MakeFuture<TProducerCommitResponse>(MockWrite(data)); + } + + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TProducerSeqNo seqNo, TData data) noexcept override { + promise.SetValue(MockWriteWithSeqNo(seqNo, data)); + } + + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TData data) noexcept override { + promise.SetValue(MockWrite(data)); + } + + NThreading::TFuture<TError> IsDead() noexcept override { + return MockIsDead(); + } + + void Cancel() override { + } + + MOCK_METHOD(TProducerCreateResponse, MockStart, (), ()); + MOCK_METHOD(NThreading::TFuture<TError>, MockIsDead, (), ()); + MOCK_METHOD(TProducerCommitResponse, MockWriteWithSeqNo, (TProducerSeqNo, TData), ()); + MOCK_METHOD(TProducerCommitResponse, MockWrite, (TData), ()); +}; + +template <class TMock = TMockProducer> +struct TCompressingProducerBootstrap { + ~TCompressingProducerBootstrap() { + Lib->CancelObjectsAndWait(); + } + + void Create(ECodec codec = ECodec::GZIP) { + Lib = new TPQLibPrivate(Settings); + MockProducer = std::make_shared<TMock>(); + CompressingProducer = std::make_shared<TLocalProducerImplCaller<TCompressingProducer>>(MockProducer, codec, 3, Lib->GetSelfRefsAreDeadPtr(), Lib, Logger); + Lib->AddToDestroySet(CompressingProducer); + } + + void MakeOKStartResponse() { + TWriteResponse wresp; + wresp.MutableInit()->SetTopic("trololo"); + TProducerCreateResponse resp(std::move(wresp)); + EXPECT_CALL(*MockProducer, MockStart()) + .WillOnce(Return(resp)); + } + + void ExpectIsDeadCall() { + EXPECT_CALL(*MockProducer, MockIsDead()) + .WillOnce(Return(DeadPromise.GetFuture())); + } + + void Start() { + MakeOKStartResponse(); + UNIT_ASSERT_STRINGS_EQUAL(CompressingProducer->Start().GetValueSync().Response.GetInit().GetTopic(), "trololo"); + } + + TPQLibSettings Settings; + TIntrusivePtr<TCerrLogger> Logger = new TCerrLogger(TLOG_DEBUG); + TIntrusivePtr<TPQLibPrivate> Lib; + std::shared_ptr<TMock> MockProducer; + std::shared_ptr<IProducerImpl> CompressingProducer; + NThreading::TPromise<TError> DeadPromise = NThreading::NewPromise<TError>(); +}; + +Y_UNIT_TEST_SUITE(TCompressingProducerTest) { + Y_UNIT_TEST(PassesStartAndDead) { + TCompressingProducerBootstrap<> bootstrap; + bootstrap.Create(); + bootstrap.ExpectIsDeadCall(); + + TWriteResponse wresp; + wresp.MutableError()->SetDescription("trololo"); + TProducerCreateResponse resp(std::move(wresp)); + EXPECT_CALL(*bootstrap.MockProducer, MockStart()) + .WillOnce(Return(resp)); + + auto startFuture = bootstrap.CompressingProducer->Start(); + UNIT_ASSERT_STRINGS_EQUAL(startFuture.GetValueSync().Response.GetError().GetDescription(), "trololo"); + + auto isDeadFuture = bootstrap.CompressingProducer->IsDead(); + UNIT_ASSERT(isDeadFuture.Initialized()); + UNIT_ASSERT(!isDeadFuture.HasValue()); + + TError err; + err.SetDescription("42"); + bootstrap.DeadPromise.SetValue(err); + + UNIT_ASSERT_STRINGS_EQUAL(isDeadFuture.GetValueSync().GetDescription(), "42"); + } + + Y_UNIT_TEST(CallsWritesInProperOrder) { + TCompressingProducerBootstrap<> bootstrap; + bootstrap.Create(); + bootstrap.Start(); + + TData data1 = TData("data1"); + TData expectedData1 = TData::Encode(data1, ECodec::GZIP, 3); + + TData data100 = TData::Raw("data100"); + TData expectedData100 = data100; + + TData data101 = TData("data101", ECodec::LZOP); + TData expectedData101 = TData::Encode(data101, ECodec::GZIP, 3); + + TWriteResponse wresp; + wresp.MutableAck(); + + InSequence seq; + EXPECT_CALL(*bootstrap.MockProducer, MockWriteWithSeqNo(1, expectedData1)) + .WillOnce(Return(TProducerCommitResponse(1, expectedData1, TWriteResponse(wresp)))); + + EXPECT_CALL(*bootstrap.MockProducer, MockWriteWithSeqNo(100, expectedData100)) + .WillOnce(Return(TProducerCommitResponse(100, expectedData100, TWriteResponse(wresp)))); + + EXPECT_CALL(*bootstrap.MockProducer, MockWrite(expectedData101)) + .WillOnce(Return(TProducerCommitResponse(101, expectedData101, TWriteResponse(wresp)))); + + auto write1 = bootstrap.CompressingProducer->Write(1, data1); + auto write100 = bootstrap.CompressingProducer->Write(100, data100); + auto write101 = bootstrap.CompressingProducer->Write(data101); + + write1.GetValueSync(); + write100.GetValueSync(); + write101.GetValueSync(); + } + + class TMockProducerWithReorder: public TMockProducer { + public: + using TMockProducer::TMockProducer; + using TMockProducer::Write; + + NThreading::TFuture<TProducerCommitResponse> Write(TProducerSeqNo seqNo, TData data) noexcept override { + return PromiseWriteWithSeqNo(seqNo, data); + } + + MOCK_METHOD(NThreading::TFuture<TProducerCommitResponse>, PromiseWriteWithSeqNo, (TProducerSeqNo, TData), ()); + }; + + void SignalsAllFuturesInProperOrderImpl(bool death) { + size_t lastWritten = 0; + + TCompressingProducerBootstrap<TMockProducerWithReorder> bootstrap; + bootstrap.Create(); + bootstrap.Start(); + + const size_t count = 200; + NThreading::TFuture<TProducerCommitResponse> responses[count]; + TData datas[count]; + + TWriteResponse wresp; + wresp.MutableAck(); + EXPECT_CALL(*bootstrap.MockProducer, PromiseWriteWithSeqNo(_, _)) + .Times(AtLeast(0)) + .WillRepeatedly(Return(NThreading::MakeFuture<TProducerCommitResponse>(TProducerCommitResponse(1, TData("data"), TWriteResponse(wresp))))); + + NThreading::TPromise<TProducerCommitResponse> promise = NThreading::NewPromise<TProducerCommitResponse>(); + + if (death) { + EXPECT_CALL(*bootstrap.MockProducer, PromiseWriteWithSeqNo(50, _)) + .Times(AtLeast(0)) + .WillRepeatedly(Return(promise.GetFuture())); + } else { + EXPECT_CALL(*bootstrap.MockProducer, PromiseWriteWithSeqNo(50, _)) + .Times(1) + .WillRepeatedly(Return(promise.GetFuture())); + } + + TString bigString(10000, 'A'); + for (size_t i = 0; i < count; ++i) { + TData data(TStringBuilder() << bigString << i, ECodec::LZOP); + datas[i] = data; + } + + auto testThreadId = TThread::CurrentThreadId(), prevThreadId = 0ul; + for (size_t i = 0; i < count; ++i) { + size_t seqNo = i + 1; + responses[i] = bootstrap.CompressingProducer->Write(i + 1, datas[i]); + auto onSignal = [&lastWritten, &prevThreadId, seqNo, testThreadId](const auto&) { + // proper thread + auto curId = TThread::CurrentThreadId(); + if (prevThreadId != 0) { + UNIT_ASSERT_C(curId == prevThreadId || curId == testThreadId, // could be executed in unittest thread if future ia already signalled or in object thread + "prevThreadId: " << Hex(prevThreadId) << ", curId: " << Hex(curId) << ", testThreadId: " << Hex(testThreadId)); + } else if (curId != testThreadId) { + prevThreadId = curId; + } + + // proper order + UNIT_ASSERT_VALUES_EQUAL(seqNo, lastWritten + 1); + lastWritten = seqNo; + }; + responses[i].Subscribe(onSignal); + } + + if (death) { + bootstrap.CompressingProducer = nullptr; + } + + responses[45].GetValueSync(); + promise.SetValue(TProducerCommitResponse(49, TData("data"), TWriteResponse(wresp))); + responses[count - 1].GetValueSync(); + for (size_t i = count - 1; i > 0; --i) { + UNIT_ASSERT(responses[i - 1].HasValue()); + } + } + + Y_UNIT_TEST(SignalsAllFuturesInProperOrderDuringDeath) { + SignalsAllFuturesInProperOrderImpl(true); + } + + Y_UNIT_TEST(SignalsAllFuturesInProperOrder) { + SignalsAllFuturesInProperOrderImpl(false); + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/consumer.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/consumer.cpp new file mode 100644 index 0000000000..da5536ab6f --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/consumer.cpp @@ -0,0 +1,612 @@ +#include "consumer.h" +#include "channel.h" +#include "persqueue_p.h" + +#include <kikimr/yndx/persqueue/read_batch_converter/read_batch_converter.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> + +#include <util/generic/strbuf.h> +#include <util/string/cast.h> +#include <util/string/printf.h> +#include <util/string/vector.h> +#include <util/string/builder.h> + +#include <grpc++/create_channel.h> + +namespace NPersQueue { + +class TConsumerDestroyHandler : public IHandler { +public: + TConsumerDestroyHandler(std::weak_ptr<TConsumer> ptr, TIntrusivePtr<TConsumer::TRpcStuff> rpcStuff, + const void* queueTag, TPQLibPrivate* pqLib) + : Ptr(std::move(ptr)) + , RpcStuff(std::move(rpcStuff)) + , QueueTag(queueTag) + , PQLib(pqLib) + {} + + void Destroy(const TError& reason) override { + auto consumer = Ptr; + auto handler = [consumer, reason] { + auto consumerShared = consumer.lock(); + if (consumerShared) { + consumerShared->Destroy(reason); + } + }; + Y_VERIFY(PQLib->GetQueuePool().GetQueue(QueueTag).AddFunc(handler)); + } + + TString ToString() override { return "DestroyHandler"; } + +protected: + std::weak_ptr<TConsumer> Ptr; + TIntrusivePtr<TConsumer::TRpcStuff> RpcStuff; // must simply live + const void* QueueTag; + TPQLibPrivate* PQLib; +}; + +class TConsumerStreamCreated : public TConsumerDestroyHandler { +public: + TConsumerStreamCreated(std::weak_ptr<TConsumer> ptr, TIntrusivePtr<TConsumer::TRpcStuff> rpcStuff, + const void* queueTag, TPQLibPrivate* pqLib) + : TConsumerDestroyHandler(std::move(ptr), std::move(rpcStuff), queueTag, pqLib) + {} + + void Done() override { + auto consumer = Ptr; + auto handler = [consumer] { + auto consumerShared = consumer.lock(); + if (consumerShared) { + consumerShared->OnStreamCreated(); + } + }; + Y_VERIFY(PQLib->GetQueuePool().GetQueue(QueueTag).AddFunc(handler)); + } + + TString ToString() override { return "StreamCreated"; } +}; + +class TConsumerWriteDone : public TConsumerDestroyHandler { +public: + TConsumerWriteDone(std::weak_ptr<TConsumer> ptr, TIntrusivePtr<TConsumer::TRpcStuff> rpcStuff, + const void* queueTag, TPQLibPrivate* pqLib) + : TConsumerDestroyHandler(std::move(ptr), std::move(rpcStuff), queueTag, pqLib) + {} + + void Done() override { + auto consumer = Ptr; + auto handler = [consumer] { + auto consumerShared = consumer.lock(); + if (consumerShared) { + consumerShared->OnWriteDone(); + } + }; + Y_VERIFY(PQLib->GetQueuePool().GetQueue(QueueTag).AddFunc(handler)); + } + + TString ToString() override { return "WriteDone"; } +}; + +class TConsumerReadDone : public TConsumerDestroyHandler { +public: + TConsumerReadDone(std::weak_ptr<TConsumer> ptr, TIntrusivePtr<TConsumer::TRpcStuff> rpcStuff, + const void* queueTag, TPQLibPrivate* pqLib) + : TConsumerDestroyHandler(std::move(ptr), std::move(rpcStuff), queueTag, pqLib) + {} + + void Done() override { + auto consumer = Ptr; + auto handler = [consumer] { + auto consumerShared = consumer.lock(); + if (consumerShared) { + consumerShared->OnReadDone(); + } + }; + Y_VERIFY(PQLib->GetQueuePool().GetQueue(QueueTag).AddFunc(handler)); + } + + TString ToString() override { return "ReadDone"; } +}; + +void TConsumer::Destroy() noexcept { + IsDestroying = true; + RpcStuff->Context.TryCancel(); + ChannelHolder.ChannelPtr = nullptr; //if waiting for channel creation + Destroy("Destructor called"); +} + +TConsumer::TConsumer(const TConsumerSettings& settings, std::shared_ptr<grpc::CompletionQueue> cq, + NThreading::TPromise<TConsumerCreateResponse> promise, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger) noexcept + : IConsumerImpl(std::move(destroyEventRef), std::move(pqLib)) + , RpcStuff(new TRpcStuff()) + , Settings(settings) + , StartPromise(promise) + , IsDeadPromise(NThreading::NewPromise<TError>()) + , UncommittedCount(0) + , UncommittedSize(0) + , MemoryUsage(0) + , ReadsOrdered(0) + , EstimateReadSize(settings.MaxSize) + , WriteInflight(true) + , ProxyCookie(0) + , Logger(std::move(logger)) + , IsDestroyed(false) + , IsDestroying(false) +{ + RpcStuff->CQ = std::move(cq); +} + +void TConsumer::Init() +{ + std::weak_ptr<TConsumer> self(shared_from_this()); + StreamCreatedHandler.Reset(new TConsumerStreamCreated(self, RpcStuff, this, PQLib.Get())); + ReadDoneHandler.Reset(new TConsumerReadDone(self, RpcStuff, this, PQLib.Get())); + WriteDoneHandler.Reset(new TConsumerWriteDone(self, RpcStuff, this, PQLib.Get())); +} + +void TConsumer::DoStart(TInstant deadline) +{ + if (IsDestroyed) + return; + + if (deadline != TInstant::Max()) { + std::weak_ptr<TConsumer> self = shared_from_this(); + auto onStartTimeout = [self] { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->OnStartTimeout(); + } + }; + StartDeadlineCallback = + PQLib->GetScheduler().Schedule(deadline, this, onStartTimeout); + } + + FillMetaHeaders(RpcStuff->Context, Settings.Server.Database, Settings.CredentialsProvider.get()); + RpcStuff->Stub = PersQueueService::NewStub(RpcStuff->Channel); + RpcStuff->Stream = RpcStuff->Stub->AsyncReadSession(&RpcStuff->Context, RpcStuff->CQ.get(), new TQueueEvent(StreamCreatedHandler)); +} + +NThreading::TFuture<TConsumerCreateResponse> TConsumer::Start(TInstant deadline) noexcept +{ + NThreading::TFuture<TConsumerCreateResponse> ret = StartPromise.GetFuture(); + if (ChannelHolder.ChannelInfo.Initialized()) { + std::weak_ptr<TConsumer> self = shared_from_this(); + PQLib->Subscribe(ChannelHolder.ChannelInfo, + this, + [self, deadline](const auto& f) { + auto sharedSelf = self.lock(); + if (sharedSelf) { + sharedSelf->SetChannel(f.GetValue()); + sharedSelf->DoStart(deadline); + } + }); + } else { + if (RpcStuff->Channel && ProxyCookie) { + DoStart(deadline); + } else { + Destroy("No channel"); + } + } + return ret; +} + +void TConsumer::SetChannel(const TChannelHolder& channel) noexcept { + if (IsDestroyed) + return; + + ChannelHolder = channel; +} + +void TConsumer::SetChannel(const TChannelInfo& channel) noexcept { + if (!channel.Channel || channel.ProxyCookie == 0) { + Destroy("Channel creation error"); + return; + } + + if (IsDestroyed) + return; + Y_VERIFY(!RpcStuff->Channel); + RpcStuff->Channel = channel.Channel; + ProxyCookie = channel.ProxyCookie; +} + + +TConsumer::~TConsumer() noexcept { + Destroy(); +} + + +/************************************************read*****************************************************/ + +void TConsumer::OrderRead() noexcept { + if (IsDeadPromise.HasValue()) { + return; + } + + if (StartPromise.Initialized()) { + return; + } + + if (MemoryUsage >= Settings.MaxMemoryUsage || ReadsOrdered >= Settings.MaxInflyRequests) { + return; + } + if (Settings.MaxUncommittedCount > 0 && UncommittedCount >= Settings.MaxUncommittedCount) { + return; + } + if (Settings.MaxUncommittedSize > 0 && UncommittedSize >= Settings.MaxUncommittedSize) { + return; + } + + + if (IsDeadPromise.HasValue()) { + return; + } else { + while (ReadsOrdered < Settings.MaxInflyRequests && MemoryUsage + ((ui64)ReadsOrdered) * EstimateReadSize <= Settings.MaxMemoryUsage) { + TReadRequest request; + Settings.CredentialsProvider->FillAuthInfo(request.MutableCredentials()); + request.MutableRead()->SetMaxCount(Settings.MaxCount); + request.MutableRead()->SetMaxSize(Settings.MaxSize); + request.MutableRead()->SetPartitionsAtOnce(Settings.PartsAtOnce); + request.MutableRead()->SetMaxTimeLagMs(Settings.MaxTimeLagMs); + request.MutableRead()->SetReadTimestampMs(Settings.ReadTimestampMs); + Requests.emplace_back(std::move(request)); + ++ReadsOrdered; + DEBUG_LOG("read ordered memusage " << MemoryUsage << " infly " << ReadsOrdered << " maxmemusage " + << Settings.MaxMemoryUsage, "", SessionId); + } + + DoWrite(); + } +} + +void TConsumer::OnWriteDone() { + WriteInflight = false; + DoWrite(); +} + +void TConsumer::DoWrite() { + if (IsDestroyed) { + return; + } + + if (WriteInflight || Requests.empty()) + return; + + WriteInflight = true; + TReadRequest req = Requests.front(); + Requests.pop_front(); + DEBUG_LOG("sending request " << req.GetRequestCase(), "", SessionId); + RpcStuff->Stream->Write(req, new TQueueEvent(WriteDoneHandler)); +} + + +void TConsumer::OnStreamCreated() { + if (IsDestroyed) { + return; + } + + TReadRequest req; + Settings.CredentialsProvider->FillAuthInfo(req.MutableCredentials()); + + TReadRequest::TInit* const init = req.MutableInit(); + init->SetClientId(Settings.ClientId); + for (const auto& t : Settings.Topics) { + init->AddTopics(t); + } + init->SetReadOnlyLocal(!Settings.ReadMirroredPartitions); + init->SetClientsideLocksAllowed(Settings.UseLockSession); + init->SetProxyCookie(ProxyCookie); + init->SetProtocolVersion(TReadRequest::ReadParamsInInit); + init->SetBalancePartitionRightNow(Settings.BalanceRightNow); + init->SetCommitsDisabled(Settings.CommitsDisabled); + + for (auto g : Settings.PartitionGroups) { + init->AddPartitionGroups(g); + } + + // Read settings + init->SetMaxReadMessagesCount(Settings.MaxCount); + init->SetMaxReadSize(Settings.MaxSize); + init->SetMaxReadPartitionsCount(Settings.PartsAtOnce); + init->SetMaxTimeLagMs(Settings.MaxTimeLagMs); + init->SetReadTimestampMs(Settings.ReadTimestampMs); + + WriteInflight = true; + RpcStuff->Stream->Write(req, new TQueueEvent(WriteDoneHandler)); + + RpcStuff->Response.Clear(); + RpcStuff->Stream->Read(&RpcStuff->Response, new TQueueEvent(ReadDoneHandler)); +} + +void TConsumer::OnReadDone() { + if (RpcStuff->Response.HasError()) { + Destroy(RpcStuff->Response.GetError()); + return; + } + + if (IsDeadPromise.HasValue()) + return; + if (StartPromise.Initialized()) { //init response + NThreading::TPromise<TConsumerCreateResponse> tmp; + tmp.Swap(StartPromise); + + Y_VERIFY(RpcStuff->Response.HasInit()); + auto res(std::move(RpcStuff->Response)); + RpcStuff->Response.Clear(); + RpcStuff->Stream->Read(&RpcStuff->Response, new TQueueEvent(ReadDoneHandler)); + + SessionId = res.GetInit().GetSessionId(); + DEBUG_LOG("read stream created", "", SessionId); + + tmp.SetValue(TConsumerCreateResponse{std::move(res)}); + + OrderRead(); + + return; + } + if (RpcStuff->Response.HasBatchedData()) { + ConvertToOldBatch(RpcStuff->Response); // LOGBROKER-3173 + } + if (RpcStuff->Response.HasData()) { //read response + NThreading::TPromise<TConsumerMessage> p(NThreading::NewPromise<TConsumerMessage>()); + MessageResponses.push_back(p.GetFuture()); + const ui32 sz = RpcStuff->Response.ByteSize(); + MemoryUsage += sz; + Y_VERIFY(ReadsOrdered); + --ReadsOrdered; + + ui32 cnt = 0; + for (ui32 i = 0; i < RpcStuff->Response.GetData().MessageBatchSize(); ++i) { + cnt += RpcStuff->Response.GetData().GetMessageBatch(i).MessageSize(); + } + if (!Settings.CommitsDisabled) { + ReadInfo.push_back({RpcStuff->Response.GetData().GetCookie(), {cnt, sz}}); + UncommittedSize += sz; + UncommittedCount += cnt; + } + + DEBUG_LOG("read done memusage " << MemoryUsage << " infly " << ReadsOrdered << " size " << sz + << " ucs " << UncommittedSize << " ucc " << UncommittedCount, "", SessionId); + + auto res(std::move(RpcStuff->Response)); + + RpcStuff->Response.Clear(); + RpcStuff->Stream->Read(&RpcStuff->Response, new TQueueEvent(ReadDoneHandler)); + + OrderRead(); + + Y_VERIFY(MemoryUsage >= sz); + EstimateReadSize = res.ByteSize(); + OrderRead(); + + p.SetValue(TConsumerMessage{std::move(res)}); + } else if (RpcStuff->Response.HasLock() || RpcStuff->Response.HasRelease() || RpcStuff->Response.HasCommit() || RpcStuff->Response.HasPartitionStatus()) { //lock/release/commit response + if (!RpcStuff->Response.HasLock()) { + INFO_LOG("got from server " << RpcStuff->Response, "", SessionId); + if (RpcStuff->Response.HasCommit() && !Settings.CommitsDisabled) { + for (ui32 i = 0; i < RpcStuff->Response.GetCommit().CookieSize(); ++i) { + ui64 cookie = RpcStuff->Response.GetCommit().GetCookie(i); + Y_VERIFY(!ReadInfo.empty() && ReadInfo.front().first == cookie); + Y_VERIFY(UncommittedCount >= ReadInfo.front().second.first); + Y_VERIFY(UncommittedSize >= ReadInfo.front().second.second); + UncommittedCount -= ReadInfo.front().second.first; + UncommittedSize -= ReadInfo.front().second.second; + ReadInfo.pop_front(); + } + } + MessageResponses.push_back(NThreading::MakeFuture<TConsumerMessage>(TConsumerMessage(std::move(RpcStuff->Response)))); + OrderRead(); + } else { + const auto& p = RpcStuff->Response.GetLock(); + NThreading::TPromise<TLockInfo> promise(NThreading::NewPromise<TLockInfo>()); + //TODO: add subscribe + auto f = promise.GetFuture(); + std::weak_ptr<TConsumer> ptr = shared_from_this(); + + //will take ptr for as long as promises are not set + PQLib->Subscribe(f, + this, + [ptr, p](const auto& f) { + auto consumer = ptr.lock(); + if (consumer) { + consumer->Lock(p.GetTopic(), p.GetPartition(), p.GetGeneration(), f.GetValue().ReadOffset, f.GetValue().CommitOffset, + f.GetValue().VerifyReadOffset); + } + }); + + INFO_LOG("got LOCK from server " << p, "", SessionId); + MessageResponses.push_back(NThreading::MakeFuture<TConsumerMessage>(TConsumerMessage(std::move(RpcStuff->Response), std::move(promise)))); + } + + RpcStuff->Response.Clear(); + RpcStuff->Stream->Read(&RpcStuff->Response, new TQueueEvent(ReadDoneHandler)); + } else { + Y_FAIL("unsupported response %s", RpcStuff->Response.DebugString().c_str()); + } + + ProcessResponses(); +} + + +/******************************************commit****************************************************/ +void TConsumer::Commit(const TVector<ui64>& cookies) noexcept { + auto sortedCookies = cookies; + std::sort(sortedCookies.begin(), sortedCookies.end()); + + TReadRequest request; + for (const auto& cookie : sortedCookies) { + request.MutableCommit()->AddCookie(cookie); + } + + DEBUG_LOG("sending COMMIT to server " << request.GetCommit(), "", SessionId); + if (IsDestroyed || StartPromise.Initialized()) { + ERR_LOG("trying to commit " << (IsDestroyed ? "after destroy" : "before start"), "", SessionId); + return; + } + Requests.emplace_back(std::move(request)); + DoWrite(); +} + + +void TConsumer::RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept { + TReadRequest request; + auto status = request.MutableStatus(); + status->SetTopic(topic); + status->SetPartition(partition); + status->SetGeneration(generation); + + DEBUG_LOG("sending GET_STATUS to server " << request.GetStatus(), "", SessionId); + if (IsDestroyed || StartPromise.Initialized()) { + ERR_LOG("trying to get status " << (IsDestroyed ? "after destroy" : "before start"), "", SessionId); + return; + } + Requests.emplace_back(std::move(request)); + DoWrite(); +} + +void TConsumer::GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept { + if (IsDeadPromise.HasValue()) { + TReadResponse res; + res.MutableError()->SetDescription("consumer is dead"); + res.MutableError()->SetCode(NErrorCode::ERROR); + promise.SetValue(TConsumerMessage{std::move(res)}); + return; + } + + if (StartPromise.Initialized()) { + TReadResponse res; + res.MutableError()->SetDescription("consumer is not ready"); + res.MutableError()->SetCode(NErrorCode::ERROR); + promise.SetValue(TConsumerMessage{std::move(res)}); + return; + } + + MessagePromises.push_back(promise); + ProcessResponses(); //if alredy have an answer +} + + +void TConsumer::Lock(const TString& topic, const ui32 partition, const ui64 generation, const ui64 readOffset, + const ui64 commitOffset, const bool verifyReadOffset) noexcept { + Y_VERIFY(Settings.UseLockSession); + + TReadRequest request; + request.MutableStartRead()->SetTopic(topic); + request.MutableStartRead()->SetPartition(partition); + request.MutableStartRead()->SetReadOffset(readOffset); + request.MutableStartRead()->SetCommitOffset(commitOffset); + request.MutableStartRead()->SetVerifyReadOffset(verifyReadOffset); + request.MutableStartRead()->SetGeneration(generation); + INFO_LOG("sending START_READ to server " << request.GetStartRead(), "", SessionId); + + Settings.CredentialsProvider->FillAuthInfo(request.MutableCredentials()); + + if (IsDestroyed) + return; + Requests.emplace_back(std::move(request)); + DoWrite(); +} + + +void TConsumer::ProcessResponses() { + while(true) { + if (MessagePromises.empty() || MessageResponses.empty() || !MessageResponses.front().HasValue()) + break; + auto p(std::move(MessagePromises.front())); + MessagePromises.pop_front(); + auto res(MessageResponses.front().ExtractValue()); + MessageResponses.pop_front(); + if (res.Type == EMT_DATA) { + ui32 sz = res.Response.ByteSize(); + Y_VERIFY(MemoryUsage >= sz); + MemoryUsage -= sz; + } + + if (res.Type == EMT_DATA) { + OrderRead(); + } + + p.SetValue(TConsumerMessage{std::move(res)}); + } +} + +/***********************************************************************************/ + +NThreading::TFuture<TError> TConsumer::IsDead() noexcept { + return IsDeadPromise.GetFuture(); +} + +void TConsumer::Destroy(const TString& description) { + TError error; + error.SetDescription(description); + error.SetCode(NErrorCode::ERROR); + Destroy(error); +} + +void TConsumer::Destroy(const TError& error) { + if (!IsDestroying) { + INFO_LOG("error: " << error, "", SessionId); + } + + if (IsDestroyed) + return; + IsDestroyed = true; + + IsDeadPromise.SetValue(error); + + MessageResponses.clear(); + + NThreading::TFuture<TChannelInfo> tmp; + tmp.Swap(ChannelHolder.ChannelInfo); + ChannelHolder.ChannelPtr = nullptr; + + Error = error; + if (StartPromise.Initialized()) { + NThreading::TPromise<TConsumerCreateResponse> tmp; + tmp.Swap(StartPromise); + TReadResponse res; + res.MutableError()->CopyFrom(Error); + tmp.SetValue(TConsumerCreateResponse{std::move(res)}); + } + + while (!MessagePromises.empty()) { + auto p = MessagePromises.front(); + MessagePromises.pop_front(); + TReadResponse res; + res.MutableError()->CopyFrom(Error); + p.SetValue(TConsumerMessage{std::move(res)}); + } + + StreamCreatedHandler.Reset(); + ReadDoneHandler.Reset(); + WriteDoneHandler.Reset(); + + DestroyPQLibRef(); +} + +void TConsumer::OnStartTimeout() { + if (IsDestroyed) { + return; + } + + StartDeadlineCallback = nullptr; + if (!StartPromise.Initialized()) { + // already replied with successful start + return; + } + + TError error; + error.SetDescription("Start timeout"); + error.SetCode(NErrorCode::CREATE_TIMEOUT); + Destroy(error); +} + +void TConsumer::Cancel() { + IsDestroying = true; + RpcStuff->Context.TryCancel(); + ChannelHolder.ChannelPtr = nullptr; //if waiting for channel creation + + Destroy(GetCancelReason()); +} + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/consumer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/consumer.h new file mode 100644 index 0000000000..f2f6f81e71 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/consumer.h @@ -0,0 +1,121 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> +#include "channel.h" +#include "scheduler.h" +#include "iconsumer_p.h" +#include "internals.h" +#include <kikimr/yndx/api/grpc/persqueue.grpc.pb.h> + +#include <library/cpp/threading/future/future.h> + +#include <deque> +#include <memory> + +namespace NPersQueue { + +class TConsumer: public IConsumerImpl, public std::enable_shared_from_this<TConsumer> { +public: + friend class TPQLibPrivate; + + NThreading::TFuture<TConsumerCreateResponse> Start(TInstant deadline) noexcept override; + + void Commit(const TVector<ui64>& cookies) noexcept override; + + using IConsumerImpl::GetNextMessage; + void GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept override; + + void RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept override; + + NThreading::TFuture<TError> IsDead() noexcept override; + + TConsumer(const TConsumerSettings& settings, std::shared_ptr<grpc::CompletionQueue> cq, + NThreading::TPromise<TConsumerCreateResponse> promise, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger) noexcept; + + void Init() override; + + ~TConsumer() noexcept; + + void Destroy() noexcept; + void SetChannel(const TChannelHolder& channel) noexcept; + void SetChannel(const TChannelInfo& channel) noexcept; + + void Lock(const TString& topic, const ui32 partition, const ui64 generation, const ui64 readOffset, const ui64 commitOffset, const bool verifyReadOffset) noexcept; + void OrderRead() noexcept; + + void Cancel() override; + +public: + using TStream = grpc::ClientAsyncReaderWriterInterface<TReadRequest, TReadResponse>; + + // objects that must live after destruction of producer untill all the callbacks arrive at CompletionQueue + struct TRpcStuff: public TAtomicRefCount<TRpcStuff> { + TReadResponse Response; + std::shared_ptr<grpc::CompletionQueue> CQ; + std::shared_ptr<grpc::Channel> Channel; + std::unique_ptr<PersQueueService::Stub> Stub; + grpc::ClientContext Context; + std::unique_ptr<TStream> Stream; + }; + +protected: + friend class TConsumerStreamCreated; + friend class TConsumerReadDone; + friend class TConsumerWriteDone; + friend class TConsumerDestroyHandler; + + IHandlerPtr StreamCreatedHandler; + IHandlerPtr ReadDoneHandler; + IHandlerPtr WriteDoneHandler; + + void ProcessResponses(); + void Destroy(const TError& reason); + void Destroy(const TString& description); // the same but with Code=ERROR + void OnStreamCreated(); + void OnReadDone(); + void OnWriteDone(); + + void DoWrite(); + void DoStart(TInstant deadline); + void OnStartTimeout(); + + TIntrusivePtr<TRpcStuff> RpcStuff; + + TChannelHolder ChannelHolder; + TConsumerSettings Settings; + + TString SessionId; + + NThreading::TPromise<TConsumerCreateResponse> StartPromise; + NThreading::TPromise<TError> IsDeadPromise; + std::deque<NThreading::TPromise<TConsumerMessage>> MessagePromises; + std::deque<NThreading::TFuture<TConsumerMessage>> MessageResponses; + + std::deque<TReadRequest> Requests; + + + std::deque<std::pair<ui64, std::pair<ui32, ui64>>> ReadInfo; //cookie -> (count, size) + ui32 UncommittedCount = 0; + ui64 UncommittedSize = 0; + + ui64 MemoryUsage = 0; + ui32 ReadsOrdered = 0; + + ui64 EstimateReadSize = 0; + + bool WriteInflight; + + ui64 ProxyCookie = 0; + + TIntrusivePtr<ILogger> Logger; + + TError Error; + + bool IsDestroyed; + bool IsDestroying; + + TIntrusivePtr<TScheduler::TCallbackHandler> StartDeadlineCallback; +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/consumer_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/consumer_ut.cpp new file mode 100644 index 0000000000..dc09bd2ef1 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/consumer_ut.cpp @@ -0,0 +1,90 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h> + +#include <library/cpp/testing/unittest/registar.h> + +namespace NPersQueue { +Y_UNIT_TEST_SUITE(TConsumerTest) { + Y_UNIT_TEST(NotStartedConsumerCanBeDestructed) { + // Test that consumer doesn't hang on till shutdown + TPQLib lib; + TConsumerSettings settings; + settings.Server = TServerSetting{"localhost"}; + settings.Topics.push_back("topic"); + settings.ClientId = "client"; + lib.CreateConsumer(settings, {}, false); + } + + Y_UNIT_TEST(StartTimeout) { + if (GrpcV1EnabledByDefault()) { + return; + } + + TTestServer testServer; + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic", 1); + testServer.WaitInit("topic"); + // Test that consumer doesn't hang on till shutdown + TConsumerSettings settings; + settings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + settings.Topics.push_back("topic"); + settings.ClientId = "client"; + auto consumer = testServer.PQLib->CreateConsumer(settings, {}, false); + auto startResult = consumer->Start(TInstant::Now()); + UNIT_ASSERT_EQUAL_C(startResult.GetValueSync().Response.GetError().GetCode(), NErrorCode::CREATE_TIMEOUT, startResult.GetValueSync().Response); + + DestroyAndWait(consumer); + } + + TProducerSettings MakeProducerSettings(const TTestServer& testServer) { + TProducerSettings producerSettings; + producerSettings.ReconnectOnFailure = false; + producerSettings.Topic = "topic1"; + producerSettings.SourceId = "123"; + producerSettings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + producerSettings.Codec = ECodec::LZOP; + return producerSettings; + } + + TConsumerSettings MakeConsumerSettings(const TTestServer& testServer) { + TConsumerSettings settings; + settings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + settings.Topics.emplace_back("topic1"); + settings.ClientId = "user"; + return settings; + } + + Y_UNIT_TEST(CancelsOperationsAfterPQLibDeath) { + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(false); + testServer.AnnoyingClient->FullInit(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + + const size_t partitions = 1; + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", partitions); + testServer.WaitInit("topic1"); + + + auto consumer = testServer.PQLib->CreateConsumer(MakeConsumerSettings(testServer), testServer.PQLibSettings.DefaultLogger, false); + UNIT_ASSERT(!consumer->Start().GetValueSync().Response.HasError()); + auto isDead = consumer->IsDead(); + UNIT_ASSERT(!isDead.HasValue()); + + auto msg1 = consumer->GetNextMessage(); + auto msg2 = consumer->GetNextMessage(); + + testServer.PQLib = nullptr; + Cerr << "PQLib destroyed" << Endl; + + UNIT_ASSERT(msg1.HasValue()); + UNIT_ASSERT(msg2.HasValue()); + + UNIT_ASSERT(msg1.GetValue().Response.HasError()); + UNIT_ASSERT(msg2.GetValue().Response.HasError()); + + auto msg3 = consumer->GetNextMessage(); + UNIT_ASSERT(msg3.HasValue()); + UNIT_ASSERT(msg3.GetValue().Response.HasError()); + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/credentials_provider.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/credentials_provider.cpp new file mode 100644 index 0000000000..e3de1ba3b5 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/credentials_provider.cpp @@ -0,0 +1,349 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/credentials_provider.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include "internals.h" + +#include <library/cpp/json/json_reader.h> +#include <library/cpp/http/simple/http_client.h> +#include <library/cpp/logger/priority.h> +#include <library/cpp/tvmauth/client/facade.h> + +#include <util/generic/ptr.h> +#include <util/string/builder.h> +#include <util/system/env.h> +#include <util/system/mutex.h> +#include <library/cpp/string_utils/quote/quote.h> +#include <ydb/public/sdk/cpp/client/ydb_driver/driver.h> +#include <kikimr/public/sdk/cpp/client/iam/iam.h> + +namespace NPersQueue { + +TString GetToken(ICredentialsProvider* credentials) { + NPersQueue::TCredentials auth; + credentials->FillAuthInfo(&auth); + TString token; + switch (auth.GetCredentialsCase()) { + case NPersQueue::TCredentials::kTvmServiceTicket: + token = auth.GetTvmServiceTicket(); + break; + case NPersQueue::TCredentials::kOauthToken: + token = auth.GetOauthToken(); + break; + case NPersQueue::TCredentials::CREDENTIALS_NOT_SET: + break; + default: + Y_VERIFY(true, "Unknown Credentials case."); + } + return token; +} + +const TString& GetTvmPqServiceName() { + static const TString PQ_SERVICE_NAME = "pq"; + return PQ_SERVICE_NAME; +} + +class TInsecureCredentialsProvider : public ICredentialsProvider { + void FillAuthInfo(NPersQueue::TCredentials*) const override {} +}; + +class TLogBridge: public NTvmAuth::ILogger { +public: + TLogBridge(TIntrusivePtr<NPersQueue::ILogger> logger = nullptr) + : Logger(logger) + {} + + void Log(int lvl, const TString& msg) override { + if (Logger) { + Logger->Log(msg, "", "", lvl); + } + } + +private: + TIntrusivePtr<NPersQueue::ILogger> Logger; +}; + +class TTVMCredentialsProvider : public ICredentialsProvider { +public: + TTVMCredentialsProvider(const NTvmAuth::NTvmApi::TClientSettings& settings, const TString& alias, TIntrusivePtr<ILogger> logger = nullptr) + : Alias(alias) + { + if (!settings.GetDstAliases().contains(Alias)) { + ythrow yexception() << "alias for `" << Alias << "` must be set"; + } + Logger = MakeIntrusive<TLogBridge>(std::move(logger)); + TvmClient = std::make_shared<NTvmAuth::TTvmClient>(settings, Logger); + } + + TTVMCredentialsProvider(const TString& secret, const ui32 srcClientId, const ui32 dstClientId, const TString& alias, TIntrusivePtr<ILogger> logger = nullptr) + : TTVMCredentialsProvider(CreateSettings(secret, srcClientId, dstClientId, alias), alias, logger) + {} + + TTVMCredentialsProvider(std::shared_ptr<NTvmAuth::TTvmClient> tvmClient, const TString& alias, TIntrusivePtr<ILogger> logger = nullptr) + : TvmClient(tvmClient) + , Alias(alias) + , Logger(MakeIntrusive<TLogBridge>(std::move(logger))) + {} + + void FillAuthInfo(NPersQueue::TCredentials* authInfo) const override { + try { + auto status = TvmClient->GetStatus(); + if (status == NTvmAuth::TClientStatus::Ok) { + authInfo->SetTvmServiceTicket(TvmClient->GetServiceTicketFor(Alias)); + } else { + Logger->Error(TStringBuilder() << "Can't get ticket: " << status.GetLastError() << "\n"); + } + } catch (...) { + Logger->Error(TStringBuilder() << "Can't get ticket: " << CurrentExceptionMessage() << "\n"); + } + } + +private: + std::shared_ptr<NTvmAuth::TTvmClient> TvmClient; + TString Alias; + NTvmAuth::TLoggerPtr Logger; + + static NTvmAuth::NTvmApi::TClientSettings CreateSettings(const TString& secret, const ui32 srcClientId, const ui32 dstClientId, const TString& alias) { + NTvmAuth::NTvmApi::TClientSettings settings; + settings.SetSelfTvmId(srcClientId); + settings.EnableServiceTicketsFetchOptions(secret, {{alias, dstClientId}}); + return settings; + } +}; + +class TTVMQloudCredentialsProvider : public ICredentialsProvider { +public: + TTVMQloudCredentialsProvider(const TString& srcAlias, const TString& dstAlias, const TDuration& refreshPeriod, TIntrusivePtr<ILogger> logger, ui32 port) + : HttpClient(TSimpleHttpClient("localhost", port)) + , Request(TStringBuilder() << "/tvm/tickets?src=" << CGIEscapeRet(srcAlias) << "&dsts=" << CGIEscapeRet(dstAlias)) + , DstAlias(dstAlias) + , LastTicketUpdate(TInstant::Zero()) + , RefreshPeriod(refreshPeriod) + , Logger(std::move(logger)) + , DstId(0) + { + GetTicket(); + } + + TTVMQloudCredentialsProvider(const ui32 srcId, const ui32 dstId, const TDuration& refreshPeriod, TIntrusivePtr<ILogger> logger, ui32 port) + : HttpClient(TSimpleHttpClient("localhost", port)) + , Request(TStringBuilder() << "/tvm/tickets?src=" << srcId << "&dsts=" << dstId) + , LastTicketUpdate(TInstant::Zero()) + , RefreshPeriod(refreshPeriod) + , Logger(std::move(logger)) + , DstId(dstId) + { + GetTicket(); + } + + void FillAuthInfo(NPersQueue::TCredentials* authInfo) const override { + if (TInstant::Now() > LastTicketUpdate + RefreshPeriod) { + GetTicket(); + } + + authInfo->SetTvmServiceTicket(Ticket); + } + +private: + TSimpleHttpClient HttpClient; + TString Request; + TString DstAlias; + mutable TString Ticket; + mutable TInstant LastTicketUpdate; + TDuration RefreshPeriod; + TIntrusivePtr<ILogger> Logger; + ui32 DstId; + + void GetTicket() const { + try { + TStringStream out; + TSimpleHttpClient::THeaders headers; + headers["Authorization"] = GetEnv("QLOUD_TVM_TOKEN"); + HttpClient.DoGet(Request, &out, headers); + NJson::TJsonValue resp; + NJson::ReadJsonTree(&out, &resp, true); + TString localDstAlias = DstAlias; + Y_VERIFY(!DstAlias.empty() || DstId != 0); + Y_VERIFY(resp.GetMap().size() == 1); + if (!localDstAlias.empty()) { + if (!resp.Has(localDstAlias)) { + ythrow yexception() << "Result doesn't contain dstAlias `" << localDstAlias << "`"; + } + } + else { + for (const auto &[k, v] : resp.GetMap()) { + if (!v.Has("tvm_id") || v["tvm_id"].GetIntegerSafe() != DstId) + ythrow yexception() << "Result doesn't contain dstId `" << DstId << "`"; + localDstAlias = k; + } + } + TString ticket = resp[localDstAlias]["ticket"].GetStringSafe(); + if (ticket.empty()) { + ythrow yexception() << "Got empty ticket"; + } + Ticket = ticket; + LastTicketUpdate = TInstant::Now(); + } catch (...) { + if (Logger) { + Logger->Log(TStringBuilder() << "Can't get ticket: " << CurrentExceptionMessage() << "\n", "", "", TLOG_ERR); + } + } + } +}; + +class TOAuthCredentialsProvider : public ICredentialsProvider { +public: + TOAuthCredentialsProvider(const TString& token) + : Token(token) + {} + + void FillAuthInfo(NPersQueue::TCredentials* authInfo) const override { + authInfo->SetOauthToken(Token); + } + +private: + TString Token; +}; + +class TTVMCredentialsForwarder : public ITVMCredentialsForwarder { +public: + TTVMCredentialsForwarder() = default; + + void FillAuthInfo(NPersQueue::TCredentials* authInfo) const override { + TGuard<TSpinLock> guard(Lock); + if (!Ticket.empty()) + authInfo->SetTvmServiceTicket(Ticket); + } + + void SetTVMServiceTicket(const TString& ticket) override { + TGuard<TSpinLock> guard(Lock); + Ticket = ticket; + } + + +private: + TString Ticket; + TSpinLock Lock; +}; + +class IIAMCredentialsProviderWrapper : public ICredentialsProvider { +public: + IIAMCredentialsProviderWrapper(std::shared_ptr<NYdb::ICredentialsProviderFactory> factory, TIntrusivePtr<ILogger> logger = nullptr) + : Provider(factory->CreateProvider()) + , Logger(std::move(logger)) + , Lock() + { + } + + void FillAuthInfo(NPersQueue::TCredentials* authInfo) const override { + TString ticket; + try { + with_lock(Lock) { + ticket = Provider->GetAuthInfo(); + } + } catch (...) { + if (Logger) { + Logger->Log(TStringBuilder() << "Can't get ticket: " << CurrentExceptionMessage() << "\n", "", "", TLOG_ERR); + } + return; + } + + authInfo->SetTvmServiceTicket(ticket); + } + +private: + std::shared_ptr<NYdb::ICredentialsProvider> Provider; + TIntrusivePtr<ILogger> Logger; + TMutex Lock; +}; + +class TIAMCredentialsProviderWrapper : public IIAMCredentialsProviderWrapper { +public: + TIAMCredentialsProviderWrapper(TIntrusivePtr<ILogger> logger = nullptr) + : IIAMCredentialsProviderWrapper( + NYdb::CreateIamCredentialsProviderFactory(), + std::move(logger) + ) + { + } +}; + +class TIAMJwtFileCredentialsProviderWrapper : public IIAMCredentialsProviderWrapper { +public: + TIAMJwtFileCredentialsProviderWrapper( + const TString& jwtKeyFilename, const TString& endpoint, const TDuration& refreshPeriod, + const TDuration& requestTimeout, TIntrusivePtr<ILogger> logger = nullptr + ) + : IIAMCredentialsProviderWrapper( + NYdb::CreateIamJwtFileCredentialsProviderFactory( + {{endpoint, refreshPeriod, requestTimeout}, jwtKeyFilename}), + std::move(logger) + ) + {} +}; + +class TIAMJwtParamsCredentialsProviderWrapper : public IIAMCredentialsProviderWrapper { +public: + TIAMJwtParamsCredentialsProviderWrapper( + const TString& jwtParams, const TString& endpoint, const TDuration& refreshPeriod, + const TDuration& requestTimeout, TIntrusivePtr<ILogger> logger = nullptr + ) + : IIAMCredentialsProviderWrapper( + NYdb::CreateIamJwtParamsCredentialsProviderFactory( + {{endpoint, refreshPeriod, requestTimeout}, jwtParams}), + std::move(logger) + ) + {} +}; + +std::shared_ptr<ICredentialsProvider> CreateInsecureCredentialsProvider() { + return std::make_shared<TInsecureCredentialsProvider>(); +} + +std::shared_ptr<ICredentialsProvider> CreateTVMCredentialsProvider(const TString& secret, const ui32 srcClientId, const ui32 dstClientId, TIntrusivePtr<ILogger> logger) { + return std::make_shared<TTVMCredentialsProvider>(secret, srcClientId, dstClientId, GetTvmPqServiceName(), std::move(logger)); +} + +std::shared_ptr<ICredentialsProvider> CreateTVMCredentialsProvider(const NTvmAuth::NTvmApi::TClientSettings& settings, TIntrusivePtr<ILogger> logger, const TString& alias) { + return std::make_shared<TTVMCredentialsProvider>(settings, alias, std::move(logger)); +} + +std::shared_ptr<ICredentialsProvider> CreateTVMCredentialsProvider(std::shared_ptr<NTvmAuth::TTvmClient> tvmClient, TIntrusivePtr<ILogger> logger, const TString& alias) { + return std::make_shared<TTVMCredentialsProvider>(tvmClient, alias, std::move(logger)); +} + +std::shared_ptr<ICredentialsProvider> CreateTVMQloudCredentialsProvider(const TString& srcAlias, const TString& dstAlias, TIntrusivePtr<ILogger> logger, const TDuration refreshPeriod, ui32 port) { + return std::make_shared<TTVMQloudCredentialsProvider>(srcAlias, dstAlias, refreshPeriod, std::move(logger), port); +} + +std::shared_ptr<ICredentialsProvider> CreateTVMQloudCredentialsProvider(const ui32 srcId, const ui32 dstId, TIntrusivePtr<ILogger> logger, const TDuration refreshPeriod, ui32 port) { + return std::make_shared<TTVMQloudCredentialsProvider>(srcId, dstId, refreshPeriod, std::move(logger), port); +} + +std::shared_ptr<ICredentialsProvider> CreateOAuthCredentialsProvider(const TString& token) { + return std::make_shared<TOAuthCredentialsProvider>(token); +} + +std::shared_ptr<ITVMCredentialsForwarder> CreateTVMCredentialsForwarder() { + return std::make_shared<TTVMCredentialsForwarder>(); +} + +std::shared_ptr<ICredentialsProvider> CreateIAMCredentialsForwarder(TIntrusivePtr<ILogger> logger) { + return std::make_shared<TIAMCredentialsProviderWrapper>(logger); +} + +std::shared_ptr<ICredentialsProvider> CreateIAMJwtFileCredentialsForwarder( + const TString& jwtKeyFilename, TIntrusivePtr<ILogger> logger, + const TString& endpoint, const TDuration& refreshPeriod, const TDuration& requestTimeout +) { + return std::make_shared<TIAMJwtFileCredentialsProviderWrapper>(jwtKeyFilename, endpoint, + refreshPeriod, requestTimeout, logger); +} + +std::shared_ptr<ICredentialsProvider> CreateIAMJwtParamsCredentialsForwarder( + const TString& jwtParams, TIntrusivePtr<ILogger> logger, + const TString& endpoint, const TDuration& refreshPeriod, const TDuration& requestTimeout +) { + return std::make_shared<TIAMJwtParamsCredentialsProviderWrapper>(jwtParams, endpoint, + refreshPeriod, requestTimeout, logger); +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/decompressing_consumer.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/decompressing_consumer.cpp new file mode 100644 index 0000000000..0fbf62cf02 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/decompressing_consumer.cpp @@ -0,0 +1,290 @@ +#include "decompressing_consumer.h" +#include "persqueue_p.h" + +#include <library/cpp/streams/lzop/lzop.h> +#include <library/cpp/streams/zstd/zstd.h> + +#include <util/stream/mem.h> +#include <util/stream/zlib.h> + +#include <atomic> + +namespace NPersQueue { + +TDecompressingConsumer::TDecompressingConsumer(std::shared_ptr<IConsumerImpl> subconsumer, const TConsumerSettings& settings, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger) + : IConsumerImpl(std::move(destroyEventRef), std::move(pqLib)) + , Logger(std::move(logger)) + , Subconsumer(std::move(subconsumer)) + , Settings(settings) +{ +} + +void TDecompressingConsumer::Init() { + DestroyedFuture = WaitExceptionOrAll(DestroyedPromise.GetFuture(), Subconsumer->Destroyed()); +} + +TDecompressingConsumer::~TDecompressingConsumer() { + DestroyQueue("Destructor called"); +} + +template <class T> +void TDecompressingConsumer::SubscribeForQueueProcessing(NThreading::TFuture<T>& future) { + const void* queueTag = this; + std::weak_ptr<TDecompressingConsumer> self = shared_from_this(); + auto signalProcessQueue = [queueTag, self, pqLib = PQLib.Get()](const auto&) mutable { + SignalProcessQueue(queueTag, std::move(self), pqLib); + }; + future.Subscribe(signalProcessQueue); +} + +NThreading::TFuture<TConsumerCreateResponse> TDecompressingConsumer::Start(TInstant deadline) noexcept { + auto ret = Subconsumer->Start(deadline); + SubscribeForQueueProcessing(ret); // when subconsumer starts, we will read ahead some messages + + // subscribe to death + std::weak_ptr<TDecompressingConsumer> self = shared_from_this(); + auto isDeadHandler = [self](const auto& error) { + auto selfShared = self.lock(); + if (selfShared && !selfShared->IsDestroyed) { + selfShared->DestroyQueue(error.GetValue()); + } + }; + PQLib->Subscribe(Subconsumer->IsDead(), this, isDeadHandler); + + return ret; +} + +NThreading::TFuture<TError> TDecompressingConsumer::IsDead() noexcept { + return IsDeadPromise.GetFuture(); +} + +void TDecompressingConsumer::Commit(const TVector<ui64>& cookies) noexcept { + if (IsDestroyed) { + ERR_LOG("Attempt to commit to dead consumer.", "", ""); + return; + } + return Subconsumer->Commit(cookies); +} + +void TDecompressingConsumer::RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept { + if (IsDestroyed) { + ERR_LOG("Attempt to request status from dead consumer.", "", ""); + return; + } + return Subconsumer->RequestPartitionStatus(topic, partition, generation); +} + + +NThreading::TFuture<void> TDecompressingConsumer::Destroyed() noexcept { + return DestroyedFuture; +} + +void TDecompressingConsumer::AddNewGetNextMessageRequest(NThreading::TPromise<TConsumerMessage>& promise) { + Queue.emplace_back(promise); + Queue.back().Future = Subconsumer->GetNextMessage(); + SubscribeForQueueProcessing(Queue.back().Future); +} + +void TDecompressingConsumer::GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept { + if (IsDestroyed) { + TReadResponse response; + auto* err = response.MutableError(); + err->SetCode(NErrorCode::ERROR); + err->SetDescription("Consumer is dead."); + promise.SetValue(TConsumerMessage(std::move(response))); + return; + } + AddNewGetNextMessageRequest(promise); +} + +void TDecompressingConsumer::ProcessQueue() { + try { + // stage 1: wait subconsumer answer and run decompression tasks + for (TReadRequestInfo& request : Queue) { + if (!request.Data && request.Future.HasValue()) { + request.Data.ConstructInPlace(request.Future.ExtractValue()); + if (request.Data->Type == EMT_DATA) { + RequestDecompressing(request); + } + } + } + + // stage 2: answer ready tasks to client + while (!Queue.empty()) { + TReadRequestInfo& front = Queue.front(); + if (!front.Data || front.Data->Type == EMT_DATA && !front.AllDecompressing.HasValue()) { + break; + } + if (front.Data->Type == EMT_DATA) { + CopyDataToAnswer(front); + } + + front.Promise.SetValue(TConsumerMessage(std::move(*front.Data))); + Queue.pop_front(); + } + } catch (const std::exception&) { + DestroyQueue(TStringBuilder() << "Failed to decompress data: " << CurrentExceptionMessage()); + } +} + +void TDecompressingConsumer::SignalProcessQueue(const void* queueTag, std::weak_ptr<TDecompressingConsumer> self, TPQLibPrivate* pqLib) { + auto processQueue = [self] { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->ProcessQueue(); + } + }; + Y_VERIFY(pqLib->GetQueuePool().GetQueue(queueTag).AddFunc(processQueue)); +} + +namespace { + +struct TWaitAll: public TAtomicRefCount<TWaitAll> { + TWaitAll(size_t count) + : Count(count) + { + } + + void OnResultReady() { + if (--Count == 0) { + Promise.SetValue(); + } + } + + std::atomic<size_t> Count; + NThreading::TPromise<void> Promise = NThreading::NewPromise<void>(); +}; + +THolder<IInputStream> CreateDecompressor(const ECodec codec, IInputStream* origin) { + THolder<IInputStream> result; + if (codec == ECodec::GZIP) { + result.Reset(new TZLibDecompress(origin)); + } else if (codec == ECodec::LZOP) { + result.Reset(new TLzopDecompress(origin)); + } else if (codec == ECodec::ZSTD) { + result.Reset(new TZstdDecompress(origin)); + } + return result; +} + +} // namespace + +void TDecompressingConsumer::RequestDecompressing(TReadRequestInfo& request) { + size_t futuresCount = 0; + const TReadResponse::TData& data = request.Data->Response.GetData(); + request.BatchFutures.resize(data.MessageBatchSize()); + for (size_t i = 0; i < data.MessageBatchSize(); ++i) { + const TReadResponse::TData::TMessageBatch& batch = data.GetMessageBatch(i); + request.BatchFutures.reserve(batch.MessageSize()); + for (const TReadResponse::TData::TMessage& message : batch.GetMessage()) { + if (message.GetMeta().GetCodec() != ECodec::RAW) { + ++futuresCount; + request.BatchFutures[i].push_back(RequestDecompressing(message)); + } + } + } + + if (futuresCount == 0) { + request.AllDecompressing = NThreading::MakeFuture(); // done + return; + } else { + TIntrusivePtr<TWaitAll> waiter = new TWaitAll(futuresCount); + auto handler = [waiter](const auto&) { + waiter->OnResultReady(); + }; + for (auto& batch : request.BatchFutures) { + for (auto& future : batch) { + future.Subscribe(handler); + } + } + request.AllDecompressing = waiter->Promise.GetFuture(); + } + SubscribeForQueueProcessing(request.AllDecompressing); +} + +NThreading::TFuture<TString> TDecompressingConsumer::RequestDecompressing(const TReadResponse::TData::TMessage& message) { + TString data = message.GetData(); + ECodec codec = message.GetMeta().GetCodec(); + NThreading::TPromise<TString> promise = NThreading::NewPromise<TString>(); + auto decompress = [data, codec, promise]() mutable { + Decompress(data, codec, promise); + }; + Y_VERIFY(PQLib->GetCompressionPool().AddFunc(decompress)); + return promise.GetFuture(); +} + +void TDecompressingConsumer::Decompress(const TString& data, ECodec codec, NThreading::TPromise<TString>& promise) { + try { + // TODO: Check if decompression was successfull, i.e. if 'data' is valid byte array compressed with 'codec' + TMemoryInput iss(data.data(), data.size()); + THolder<IInputStream> dec = CreateDecompressor(codec, &iss); + if (!dec) { + ythrow yexception() << "Failed to create decompressor"; + } + + TString result; + TStringOutput oss(result); + TransferData(dec.Get(), &oss); + + promise.SetValue(result); + } catch (...) { + promise.SetException(std::current_exception()); + } +} + +void TDecompressingConsumer::CopyDataToAnswer(TReadRequestInfo& request) { + TReadResponse::TData& data = *request.Data->Response.MutableData(); + Y_VERIFY(request.BatchFutures.size() == data.MessageBatchSize()); + for (size_t i = 0; i < request.BatchFutures.size(); ++i) { + TReadResponse::TData::TMessageBatch& batch = *data.MutableMessageBatch(i); + auto currentMessage = request.BatchFutures[i].begin(); + for (TReadResponse::TData::TMessage& message : *batch.MutableMessage()) { + if (message.GetMeta().GetCodec() != ECodec::RAW) { + Y_VERIFY(currentMessage != request.BatchFutures[i].end()); + try { + message.SetData(currentMessage->GetValue()); + message.MutableMeta()->SetCodec(ECodec::RAW); + } catch (const std::exception& ex) { + if (Settings.SkipBrokenChunks) { + message.SetBrokenPackedData(message.GetData()); + message.SetData(""); + } else { + throw; + } + } + ++currentMessage; + } + } + Y_VERIFY(currentMessage == request.BatchFutures[i].end()); + } +} + +void TDecompressingConsumer::DestroyQueue(const TString& errorMessage) { + TError error; + error.SetDescription(errorMessage); + error.SetCode(NErrorCode::ERROR); + DestroyQueue(error); +} + +void TDecompressingConsumer::DestroyQueue(const TError& error) { + if (IsDestroyed) { + return; + } + + for (TReadRequestInfo& request : Queue) { + TReadResponse response; + *response.MutableError() = error; + request.Promise.SetValue(TConsumerMessage(std::move(response))); + } + Queue.clear(); + IsDeadPromise.SetValue(error); + IsDestroyed = true; + + DestroyPQLibRef(); +} + +void TDecompressingConsumer::Cancel() { + DestroyQueue(GetCancelReason()); +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/decompressing_consumer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/decompressing_consumer.h new file mode 100644 index 0000000000..c547ef8b3a --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/decompressing_consumer.h @@ -0,0 +1,77 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> +#include "internals.h" +#include "iconsumer_p.h" + +#include <util/generic/maybe.h> + +#include <deque> +#include <memory> +#include <vector> + +namespace NPersQueue { + +class TPQLibPrivate; + +class TDecompressingConsumer: public IConsumerImpl, public std::enable_shared_from_this<TDecompressingConsumer> { +public: + using IConsumerImpl::GetNextMessage; + NThreading::TFuture<TConsumerCreateResponse> Start(TInstant deadline) noexcept override; + NThreading::TFuture<TError> IsDead() noexcept override; + NThreading::TFuture<void> Destroyed() noexcept override; + void GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept override; + void Commit(const TVector<ui64>& cookies) noexcept override; + void RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept override; + void Init() override; + + ~TDecompressingConsumer(); + TDecompressingConsumer(std::shared_ptr<IConsumerImpl> subconsumer, const TConsumerSettings& settings, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger); + + void Cancel() override; + +private: + struct TReadRequestInfo { + TReadRequestInfo() = default; + + TReadRequestInfo(const NThreading::TPromise<TConsumerMessage>& promise) + : Promise(promise) + { + } + + NThreading::TPromise<TConsumerMessage> Promise; + NThreading::TFuture<TConsumerMessage> Future; + TMaybe<TConsumerMessage> Data; + + using TMessageFutures = std::vector<NThreading::TFuture<TString>>; + using TBatchFutures = std::vector<TMessageFutures>; + TBatchFutures BatchFutures; + NThreading::TFuture<void> AllDecompressing; + }; + + void ProcessQueue(); + static void SignalProcessQueue(const void* queueTag, std::weak_ptr<TDecompressingConsumer> self, TPQLibPrivate* pqLib); + void DestroyQueue(const TString& errorMessage); + void DestroyQueue(const TError& error); + + template <class T> + void SubscribeForQueueProcessing(NThreading::TFuture<T>& future); + + void RequestDecompressing(TReadRequestInfo& request); + void CopyDataToAnswer(TReadRequestInfo& request); + NThreading::TFuture<TString> RequestDecompressing(const TReadResponse::TData::TMessage& message); + static void Decompress(const TString& data, ECodec codec, NThreading::TPromise<TString>& promise); + void AddNewGetNextMessageRequest(NThreading::TPromise<TConsumerMessage>& promise); + +protected: + TIntrusivePtr<ILogger> Logger; + std::shared_ptr<IConsumerImpl> Subconsumer; + std::deque<TReadRequestInfo> Queue; + TConsumerSettings Settings; + NThreading::TFuture<void> DestroyedFuture; // Subconsumer may be deleted, so store our future here. + NThreading::TPromise<TError> IsDeadPromise = NThreading::NewPromise<TError>(); + bool IsDestroyed = false; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/decompressing_consumer_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/decompressing_consumer_ut.cpp new file mode 100644 index 0000000000..6c22163e8f --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/decompressing_consumer_ut.cpp @@ -0,0 +1,265 @@ +#include "decompressing_consumer.h" +#include "local_caller.h" +#include "persqueue_p.h" +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h> + +#include <library/cpp/testing/gmock_in_unittest/gmock.h> +#include <library/cpp/testing/unittest/registar.h> + +using namespace testing; + +namespace NPersQueue { + +class TMockConsumer: public IConsumerImpl { +public: + TMockConsumer() + : IConsumerImpl(nullptr, nullptr) + { + } + + NThreading::TFuture<TConsumerCreateResponse> Start(TInstant) noexcept override { + return NThreading::MakeFuture<TConsumerCreateResponse>(MockStart()); + } + + NThreading::TFuture<TError> IsDead() noexcept override { + return MockIsDead(); + } + + NThreading::TFuture<TConsumerMessage> GetNextMessage() noexcept override { + return NThreading::MakeFuture<TConsumerMessage>(MockGetNextMessage()); + } + + void GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept override { + promise.SetValue(MockGetNextMessage()); + } + + void Commit(const TVector<ui64>&) noexcept override { + } + + void RequestPartitionStatus(const TString&, ui64, ui64) noexcept override { + } + + void Cancel() override { + } + + MOCK_METHOD(TConsumerCreateResponse, MockStart, (), ()); + MOCK_METHOD(NThreading::TFuture<TError>, MockIsDead, (), ()); + MOCK_METHOD(TConsumerMessage, MockGetNextMessage, (), ()); +}; + +template <class TMock = TMockConsumer> +struct TDecompressingConsumerBootstrap { + ~TDecompressingConsumerBootstrap() { + Lib->CancelObjectsAndWait(); + } + + void Create() { + Lib = new TPQLibPrivate(PQLibSettings); + MockConsumer = std::make_shared<TMock>(); + DecompressingConsumer = std::make_shared<TLocalConsumerImplCaller<TDecompressingConsumer>>(MockConsumer, Settings, Lib->GetSelfRefsAreDeadPtr(), Lib, Logger); + Lib->AddToDestroySet(DecompressingConsumer); + } + + void MakeOKStartResponse() { + TReadResponse rresp; + rresp.MutableInit()->SetSessionId("test-session"); + TConsumerCreateResponse resp(std::move(rresp)); + EXPECT_CALL(*MockConsumer, MockStart()) + .WillOnce(Return(resp)); + } + + void ExpectIsDeadCall() { + EXPECT_CALL(*MockConsumer, MockIsDead()) + .WillOnce(Return(DeadPromise.GetFuture())); + } + + void Start() { + MakeOKStartResponse(); + UNIT_ASSERT_STRINGS_EQUAL(DecompressingConsumer->Start().GetValueSync().Response.GetInit().GetSessionId(), "test-session"); + } + + TConsumerSettings Settings; + TPQLibSettings PQLibSettings; + TIntrusivePtr<TCerrLogger> Logger = new TCerrLogger(TLOG_DEBUG); + TIntrusivePtr<TPQLibPrivate> Lib; + std::shared_ptr<TMock> MockConsumer; + std::shared_ptr<IConsumerImpl> DecompressingConsumer; + NThreading::TPromise<TError> DeadPromise = NThreading::NewPromise<TError>(); +}; + +Y_UNIT_TEST_SUITE(TDecompressingConsumerTest) { + Y_UNIT_TEST(DiesOnDeadSubconsumer) { + TDecompressingConsumerBootstrap<> bootstrap; + bootstrap.Create(); + bootstrap.ExpectIsDeadCall(); + bootstrap.Start(); + + auto isDead = bootstrap.DecompressingConsumer->IsDead(); + UNIT_ASSERT(!isDead.HasValue()); + + bootstrap.DeadPromise.SetValue(TError()); + + isDead.GetValueSync(); // doesn't hang on + } + + Y_UNIT_TEST(PassesNonData) { + TDecompressingConsumerBootstrap<> bootstrap; + bootstrap.Create(); + bootstrap.ExpectIsDeadCall(); + bootstrap.Start(); + + InSequence sequence; + { + TReadResponse resp; + resp.MutableCommit()->AddCookie(42); + + TConsumerMessage ret(std::move(resp)); + + EXPECT_CALL(*bootstrap.MockConsumer, MockGetNextMessage()) + .WillOnce(Return(ret)); + } + + { + TReadResponse resp; + resp.MutableRelease()->SetTopic("topic!"); + + TConsumerMessage ret(std::move(resp)); + + EXPECT_CALL(*bootstrap.MockConsumer, MockGetNextMessage()) + .WillOnce(Return(ret)); + } + + auto passed1 = bootstrap.DecompressingConsumer->GetNextMessage().GetValueSync(); + UNIT_ASSERT_VALUES_EQUAL_C(passed1.Response.GetCommit().CookieSize(), 1, passed1.Response); + UNIT_ASSERT_VALUES_EQUAL_C(passed1.Response.GetCommit().GetCookie(0), 42, passed1.Response); + + auto passed2 = bootstrap.DecompressingConsumer->GetNextMessage().GetValueSync(); + UNIT_ASSERT_STRINGS_EQUAL_C(passed2.Response.GetRelease().GetTopic(), "topic!", passed2.Response); + } + + void ProcessesBrokenChunks(bool skip) { + TDecompressingConsumerBootstrap<> bootstrap; + bootstrap.Settings.SkipBrokenChunks = skip; + bootstrap.Create(); + bootstrap.ExpectIsDeadCall(); + bootstrap.Start(); + + TReadResponse resp; + auto* msg = resp.MutableData()->AddMessageBatch()->AddMessage(); + msg->SetData("hjdhkjhkjhshqsiuhqisuqihsi;"); + msg->MutableMeta()->SetCodec(ECodec::LZOP); + + EXPECT_CALL(*bootstrap.MockConsumer, MockGetNextMessage()) + .WillOnce(Return(TConsumerMessage(std::move(resp)))); + + auto isDead = bootstrap.DecompressingConsumer->IsDead(); + UNIT_ASSERT(!isDead.HasValue()); + + auto passed = bootstrap.DecompressingConsumer->GetNextMessage().GetValueSync(); + + if (skip) { + UNIT_ASSERT(!passed.Response.HasError()); + isDead.HasValue(); + UNIT_ASSERT_VALUES_EQUAL(passed.Response.GetData().MessageBatchSize(), 1); + UNIT_ASSERT_VALUES_EQUAL(passed.Response.GetData().GetMessageBatch(0).MessageSize(), 1); + UNIT_ASSERT(passed.Response.GetData().GetMessageBatch(0).GetMessage(0).GetData().empty()); + } else { + UNIT_ASSERT(passed.Response.HasError()); + isDead.GetValueSync(); + + UNIT_ASSERT(bootstrap.DecompressingConsumer->GetNextMessage().GetValueSync().Response.HasError()); + } + + //DestroyAndWait(bootstrap.DecompressingConsumer); + } + + Y_UNIT_TEST(DiesOnBrokenChunks) { + ProcessesBrokenChunks(false); + } + + Y_UNIT_TEST(SkipsBrokenChunks) { + ProcessesBrokenChunks(true); + } + + static void AddMessage(TReadResponse::TData::TMessageBatch* batch, const TString& sourceData, ECodec codec = ECodec::RAW) { + auto* msg = batch->AddMessage(); + msg->MutableMeta()->SetCodec(codec); + if (codec == ECodec::RAW) { + msg->SetData(sourceData); + } else { + msg->SetData(TData::Encode(sourceData, codec, -1).GetEncodedData()); + } + } + + Y_UNIT_TEST(DecompessesDataInProperChuncksOrder) { + TDecompressingConsumerBootstrap<> bootstrap; + bootstrap.Create(); + bootstrap.ExpectIsDeadCall(); + bootstrap.Start(); + + InSequence sequence; + { + TReadResponse resp; + auto* batch = resp.MutableData()->AddMessageBatch(); + AddMessage(batch, "message1", ECodec::LZOP); + AddMessage(batch, "message2"); + AddMessage(batch, "message3", ECodec::GZIP); + + resp.MutableData()->AddMessageBatch(); + + batch = resp.MutableData()->AddMessageBatch(); + AddMessage(batch, "messageA", ECodec::LZOP); + AddMessage(batch, "messageB", ECodec::ZSTD); + AddMessage(batch, "messageC"); + + EXPECT_CALL(*bootstrap.MockConsumer, MockGetNextMessage()) + .WillOnce(Return(TConsumerMessage(std::move(resp)))); + } + + { + TReadResponse resp; + AddMessage(resp.MutableData()->AddMessageBatch(), "trololo", ECodec::LZOP); + + EXPECT_CALL(*bootstrap.MockConsumer, MockGetNextMessage()) + .WillOnce(Return(TConsumerMessage(std::move(resp)))); + } + + auto isDead = bootstrap.DecompressingConsumer->IsDead(); + UNIT_ASSERT(!isDead.HasValue()); + + auto f1 = bootstrap.DecompressingConsumer->GetNextMessage(); + auto f2 = bootstrap.DecompressingConsumer->GetNextMessage(); + + auto data1 = f1.GetValueSync().Response.GetData(); + auto data2 = f2.GetValueSync().Response.GetData(); + + UNIT_ASSERT_VALUES_EQUAL(data1.MessageBatchSize(), 3); + UNIT_ASSERT_VALUES_EQUAL(data1.GetMessageBatch(0).MessageSize(), 3); + UNIT_ASSERT_VALUES_EQUAL(data1.GetMessageBatch(1).MessageSize(), 0); + UNIT_ASSERT_VALUES_EQUAL(data1.GetMessageBatch(2).MessageSize(), 3); + + UNIT_ASSERT_VALUES_EQUAL(data1.GetMessageBatch(0).GetMessage(0).GetData(), "message1"); + UNIT_ASSERT_EQUAL(data1.GetMessageBatch(0).GetMessage(0).GetMeta().GetCodec(), ECodec::RAW); + UNIT_ASSERT_VALUES_EQUAL(data1.GetMessageBatch(0).GetMessage(1).GetData(), "message2"); + UNIT_ASSERT_EQUAL(data1.GetMessageBatch(0).GetMessage(1).GetMeta().GetCodec(), ECodec::RAW); + UNIT_ASSERT_VALUES_EQUAL(data1.GetMessageBatch(0).GetMessage(2).GetData(), "message3"); + UNIT_ASSERT_EQUAL(data1.GetMessageBatch(0).GetMessage(2).GetMeta().GetCodec(), ECodec::RAW); + + UNIT_ASSERT_VALUES_EQUAL(data1.GetMessageBatch(2).GetMessage(0).GetData(), "messageA"); + UNIT_ASSERT_EQUAL(data1.GetMessageBatch(2).GetMessage(0).GetMeta().GetCodec(), ECodec::RAW); + UNIT_ASSERT_VALUES_EQUAL(data1.GetMessageBatch(2).GetMessage(1).GetData(), "messageB"); + UNIT_ASSERT_EQUAL(data1.GetMessageBatch(2).GetMessage(1).GetMeta().GetCodec(), ECodec::RAW); + UNIT_ASSERT_VALUES_EQUAL(data1.GetMessageBatch(2).GetMessage(2).GetData(), "messageC"); + UNIT_ASSERT_EQUAL(data1.GetMessageBatch(2).GetMessage(2).GetMeta().GetCodec(), ECodec::RAW); + + + UNIT_ASSERT_VALUES_EQUAL(data2.MessageBatchSize(), 1); + UNIT_ASSERT_VALUES_EQUAL(data2.GetMessageBatch(0).MessageSize(), 1); + UNIT_ASSERT_VALUES_EQUAL(data2.GetMessageBatch(0).GetMessage(0).GetData(), "trololo"); + UNIT_ASSERT_EQUAL(data2.GetMessageBatch(0).GetMessage(0).GetMeta().GetCodec(), ECodec::RAW); + + + UNIT_ASSERT(!isDead.HasValue()); + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iconsumer_p.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iconsumer_p.cpp new file mode 100644 index 0000000000..2ebfcd8d07 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iconsumer_p.cpp @@ -0,0 +1,20 @@ +#include "iconsumer_p.h" +#include "persqueue_p.h" + +namespace NPersQueue { + +TPublicConsumer::TPublicConsumer(std::shared_ptr<IConsumerImpl> impl) + : Impl(std::move(impl)) +{ +} + +TPublicConsumer::~TPublicConsumer() { + Impl->Cancel(); +} + +IConsumerImpl::IConsumerImpl(std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib) + : TSyncDestroyed(std::move(destroyEventRef), std::move(pqLib)) +{ +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iconsumer_p.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iconsumer_p.h new file mode 100644 index 0000000000..35c151b9de --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iconsumer_p.h @@ -0,0 +1,51 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iconsumer.h> +#include "interface_common.h" + +namespace NPersQueue { + +class TPQLibPrivate; + +class IConsumerImpl: public IConsumer, public TSyncDestroyed { +public: + IConsumerImpl(std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib); + + virtual void Init() { + } + + using IConsumer::GetNextMessage; + virtual void GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept = 0; +}; + +// Consumer that is given to client. +class TPublicConsumer: public IConsumer { +public: + explicit TPublicConsumer(std::shared_ptr<IConsumerImpl> impl); + ~TPublicConsumer(); + + NThreading::TFuture<TConsumerCreateResponse> Start(TInstant deadline) noexcept override { + return Impl->Start(deadline); + } + + NThreading::TFuture<TConsumerMessage> GetNextMessage() noexcept override { + return Impl->GetNextMessage(); + } + + void RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept override { + Impl->RequestPartitionStatus(topic, partition, generation); + } + + void Commit(const TVector<ui64>& cookies) noexcept override { + Impl->Commit(cookies); + } + + NThreading::TFuture<TError> IsDead() noexcept override { + return Impl->IsDead(); + } + +private: + std::shared_ptr<IConsumerImpl> Impl; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/interface_common.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/interface_common.cpp new file mode 100644 index 0000000000..e88c5b1280 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/interface_common.cpp @@ -0,0 +1,31 @@ +#include "interface_common.h" +#include "persqueue_p.h" + +namespace NPersQueue { + +const TString& GetCancelReason() +{ + static const TString reason = "Destroyed"; + return reason; +} + +TSyncDestroyed::TSyncDestroyed(std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib) + : DestroyEventRef(std::move(destroyEventRef)) + , PQLib(std::move(pqLib)) +{ +} + +TSyncDestroyed::~TSyncDestroyed() { + DestroyedPromise.SetValue(); +} + +void TSyncDestroyed::DestroyPQLibRef() { + auto guard = Guard(DestroyLock); + if (DestroyEventRef) { + IsCanceling = true; + PQLib = nullptr; + DestroyEventRef = nullptr; + } +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/interface_common.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/interface_common.h new file mode 100644 index 0000000000..bd3a228b2c --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/interface_common.h @@ -0,0 +1,50 @@ +#pragma once + +#include <library/cpp/threading/future/future.h> + +#include <util/generic/ptr.h> +#include <util/system/spinlock.h> + +#include <memory> + +namespace NPersQueue { + +class TPQLibPrivate; + +const TString& GetCancelReason(); + +class TSyncDestroyed { +protected: + TSyncDestroyed(std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib); + +public: + // Get future that is signalled after object's destructor. + // Non-pqlib threads can wait on this future to ensure that + // all async operations are finished. + virtual NThreading::TFuture<void> Destroyed() noexcept { + return DestroyedPromise.GetFuture(); + } + + // This destructor will be executed after real object's destructor + // in which it is expected to set all owned promises. + virtual ~TSyncDestroyed(); + + virtual void Cancel() = 0; + + void SetDestroyEventRef(std::shared_ptr<void> ref) { + DestroyEventRef = std::move(ref); + } + +protected: + // Needed for proper PQLib deinitialization + void DestroyPQLibRef(); + +protected: + NThreading::TPromise<void> DestroyedPromise = NThreading::NewPromise<void>(); + std::shared_ptr<void> DestroyEventRef; + TIntrusivePtr<TPQLibPrivate> PQLib; + bool IsCanceling = false; + TAdaptiveLock DestroyLock; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/internals.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/internals.h new file mode 100644 index 0000000000..5c893da75e --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/internals.h @@ -0,0 +1,105 @@ +#pragma once +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> + +#include <library/cpp/grpc/common/time_point.h> +#include <library/cpp/threading/future/future.h> + +#include <util/generic/ptr.h> +#include <util/generic/string.h> +#include <util/string/cast.h> + +#include <grpc++/channel.h> +#include <grpc++/create_channel.h> + +#include <contrib/libs/grpc/include/grpcpp/impl/codegen/client_context.h> + +#include <chrono> + +#define WRITE_LOG(msg, srcId, sessionId, level, logger) \ + if (logger && logger->IsEnabled(level)) { \ + logger->Log(TStringBuilder() << msg, srcId, sessionId, level); \ + } + +#define DEBUG_LOG(msg, srcId, sessionId) WRITE_LOG(msg, srcId, sessionId, TLOG_DEBUG, Logger) +#define INFO_LOG(msg, srcId, sessionId) WRITE_LOG(msg, srcId, sessionId, TLOG_INFO, Logger) +#define WARN_LOG(msg, srcId, sessionId) WRITE_LOG(msg, srcId, sessionId, TLOG_WARNING, Logger) +#define ERR_LOG(msg, srcId, sessionId) WRITE_LOG(msg, srcId, sessionId, TLOG_ERR, Logger) + +namespace NPersQueue { + +TString GetToken(ICredentialsProvider* credentials); + +void FillMetaHeaders(grpc::ClientContext& context, const TString& database, ICredentialsProvider* credentials); + +bool UseCDS(const TServerSetting& server); + +struct TWriteData { + TProducerSeqNo SeqNo; + TData Data; + + TWriteData(ui64 seqNo, TData&& data) + : SeqNo(seqNo) + , Data(std::move(data)) + {} +}; + +class IQueueEvent { +public: + virtual ~IQueueEvent() = default; + + //! Execute an action defined by implementation. + virtual bool Execute(bool ok) = 0; + + //! Finish and destroy request. + virtual void DestroyRequest() = 0; +}; + +class IHandler : public TAtomicRefCount<IHandler> { +public: + IHandler() + {} + + virtual ~IHandler() + {} + + virtual void Destroy(const TError&) = 0; + virtual void Done() = 0; + virtual TString ToString() = 0; +}; + +using IHandlerPtr = TIntrusivePtr<IHandler>; + +class TQueueEvent : public IQueueEvent { +public: + TQueueEvent(IHandlerPtr handler) + : Handler(std::move(handler)) + { + Y_ASSERT(Handler); + } + +private: + bool Execute(bool ok) override { + if (ok) { + Handler->Done(); + } else { + TError error; + error.SetDescription("event " + Handler->ToString() + " failed"); + error.SetCode(NErrorCode::ERROR); + Handler->Destroy(error); + } + return false; + } + + virtual ~TQueueEvent() { + Handler.Reset(); + } + + void DestroyRequest() override + { + delete this; + } + + IHandlerPtr Handler; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iprocessor_p.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iprocessor_p.cpp new file mode 100644 index 0000000000..f245df74fe --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iprocessor_p.cpp @@ -0,0 +1,20 @@ +#include "iprocessor_p.h" +#include "persqueue_p.h" + +namespace NPersQueue { + +TPublicProcessor::TPublicProcessor(std::shared_ptr<IProcessorImpl> impl) + : Impl(std::move(impl)) +{ +} + +TPublicProcessor::~TPublicProcessor() { + Impl->Cancel(); +} + +IProcessorImpl::IProcessorImpl(std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib) + : TSyncDestroyed(std::move(destroyEventRef), std::move(pqLib)) +{ +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iprocessor_p.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iprocessor_p.h new file mode 100644 index 0000000000..e315b8412f --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iprocessor_p.h @@ -0,0 +1,35 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iprocessor.h> +#include "interface_common.h" + +namespace NPersQueue { + +class TPQLibPrivate; + +struct IProcessorImpl: public IProcessor, public TSyncDestroyed { + using IProcessor::GetNextData; + + IProcessorImpl(std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib); + + // Initialization after constructor (for example, for correct call of shared_from_this()) + virtual void Init() { + } + virtual void GetNextData(NThreading::TPromise<TOriginData>& promise) noexcept = 0; +}; + +// Processor that is given to client. +class TPublicProcessor: public IProcessor { +public: + explicit TPublicProcessor(std::shared_ptr<IProcessorImpl> impl); + ~TPublicProcessor(); + + NThreading::TFuture<TOriginData> GetNextData() noexcept override { + return Impl->GetNextData(); + } + +private: + std::shared_ptr<IProcessorImpl> Impl; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iproducer_p.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iproducer_p.cpp new file mode 100644 index 0000000000..c7a267b862 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iproducer_p.cpp @@ -0,0 +1,20 @@ +#include "iproducer_p.h" +#include "persqueue_p.h" + +namespace NPersQueue { + +TPublicProducer::TPublicProducer(std::shared_ptr<IProducerImpl> impl) + : Impl(std::move(impl)) +{ +} + +TPublicProducer::~TPublicProducer() { + Impl->Cancel(); +} + +IProducerImpl::IProducerImpl(std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib) + : TSyncDestroyed(std::move(destroyEventRef), std::move(pqLib)) +{ +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iproducer_p.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iproducer_p.h new file mode 100644 index 0000000000..437798c551 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iproducer_p.h @@ -0,0 +1,58 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iproducer.h> +#include "interface_common.h" + +namespace NPersQueue { + +class TPQLibPrivate; + +struct IProducerImpl: public IProducer, public TSyncDestroyed { + using IProducer::Write; + + IProducerImpl(std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib); + + // Initialization after constructor (for example, for correct call of shared_from_this()) + virtual void Init() { + } + virtual void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TProducerSeqNo seqNo, TData data) noexcept = 0; + virtual void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TData data) noexcept = 0; + + NThreading::TFuture<TProducerCommitResponse> Write(TProducerSeqNo, TData) noexcept override { + Y_FAIL(""); + return NThreading::NewPromise<TProducerCommitResponse>().GetFuture(); + } + + NThreading::TFuture<TProducerCommitResponse> Write(TData) noexcept override { + Y_FAIL(""); + return NThreading::NewPromise<TProducerCommitResponse>().GetFuture(); + } +}; + +// Producer that is given to client. +class TPublicProducer: public IProducer { +public: + explicit TPublicProducer(std::shared_ptr<IProducerImpl> impl); + ~TPublicProducer(); + + NThreading::TFuture<TProducerCreateResponse> Start(TInstant deadline = TInstant::Max()) noexcept { + return Impl->Start(deadline); + } + + NThreading::TFuture<TProducerCommitResponse> Write(TProducerSeqNo seqNo, TData data) noexcept { + return Impl->Write(seqNo, std::move(data)); + } + + NThreading::TFuture<TProducerCommitResponse> Write(TData data) noexcept { + return Impl->Write(std::move(data)); + } + + NThreading::TFuture<TError> IsDead() noexcept { + return Impl->IsDead(); + } + +private: + std::shared_ptr<IProducerImpl> Impl; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/local_caller.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/local_caller.h new file mode 100644 index 0000000000..2ac0b3ec18 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/local_caller.h @@ -0,0 +1,303 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include "interface_common.h" +#include "internals.h" +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iproducer.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> +#include "persqueue_p.h" + +namespace NPersQueue { + +template <class TImpl> +class TCancelCaller: public TImpl { +public: + using TImpl::TImpl; + + std::shared_ptr<TCancelCaller> MakeSharedRef() { + return std::static_pointer_cast<TCancelCaller>(this->shared_from_this()); + } + + void Cancel() override { + auto guard = Guard(this->DestroyLock); + if (!this->IsCanceling) { + this->IsCanceling = true; + AddCancelCaller(); + } + } + +protected: + const TImpl* QueueTag() const { // we must have exactly the same thread tag as TImpl object + return this; + } + + static void MakeCanceledError(TError& error) { + error.SetDescription(GetCancelReason()); + error.SetCode(NErrorCode::ERROR); + } + + template <class TResponseProto> + static void MakeCanceledError(TResponseProto& resp) { + MakeCanceledError(*resp.MutableError()); + } + + template <class TResponse, class... TArgs> + static void MakeCanceledError(NThreading::TPromise<TResponse>& promise, TArgs&&...) { // the most common case for calls + decltype(promise.GetValue().Response) responseProto; + MakeCanceledError(responseProto); + promise.SetValue(TResponse(std::move(responseProto))); + } + + template <class TResponse> + static TResponse MakeCanceledError() { + TResponse resp; + MakeCanceledError(*resp.MutableError()); + return resp; + } + +private: + void AddCancelCaller() { + std::weak_ptr<TCancelCaller> selfWeak = MakeSharedRef(); + auto caller = [selfWeak]() mutable { + auto selfShared = selfWeak.lock(); + if (selfShared) { + selfShared->TImpl::Cancel(); + } + }; + Y_VERIFY(this->PQLib->GetQueuePool().GetQueue(QueueTag()).AddFunc(caller)); + } +}; + +template <class TImpl, class TStartResponse> +class TLocalStartDeadImplCaller: public TCancelCaller<TImpl> { +protected: + using TCancelCaller<TImpl>::TCancelCaller; + + NThreading::TFuture<TStartResponse> Start(TInstant deadline) noexcept override { + IsDead(); // subscribe for impl future + auto guard = Guard(this->DestroyLock); + if (this->IsCanceling) { + guard.Release(); + decltype(reinterpret_cast<TStartResponse*>(0)->Response) respProto; + this->MakeCanceledError(respProto); + return NThreading::MakeFuture(TStartResponse(std::move(respProto))); + } + if (StartPromise.Initialized()) { + guard.Release(); + decltype(reinterpret_cast<TStartResponse*>(0)->Response) respProto; + TError& error = *respProto.MutableError(); + error.SetDescription("Start was already called"); + error.SetCode(NErrorCode::BAD_REQUEST); + return NThreading::MakeFuture(TStartResponse(std::move(respProto))); + } + StartPromise = NThreading::NewPromise<TStartResponse>(); + AddStartCaller(deadline); + return StartPromise.GetFuture(); + } + + NThreading::TFuture<TError> IsDead() noexcept override { + auto guard = Guard(this->DestroyLock); + if (this->IsCanceling && !IsDeadCalled) { + guard.Release(); + TError err; + this->MakeCanceledError(err); + return NThreading::MakeFuture(err); + } + if (!IsDeadCalled) { + AddIsDeadCaller(); + } + return DeadPromise.GetFuture(); + } + +private: + std::shared_ptr<TLocalStartDeadImplCaller> MakeSharedRef() { + return std::static_pointer_cast<TLocalStartDeadImplCaller>(this->shared_from_this()); + } + + void AddIsDeadCaller() { + IsDeadCalled = true; + std::shared_ptr<TLocalStartDeadImplCaller> selfShared = MakeSharedRef(); + auto isDeadCaller = [selfShared]() mutable { + selfShared->CallIsDead(); + }; + Y_VERIFY(this->PQLib->GetQueuePool().GetQueue(this->QueueTag()).AddFunc(isDeadCaller)); + } + + void CallIsDead() { + auto deadPromise = DeadPromise; + TImpl::IsDead().Subscribe([deadPromise](const auto& future) mutable { + deadPromise.SetValue(future.GetValue()); + }); + } + + void AddStartCaller(TInstant deadline) { + std::shared_ptr<TLocalStartDeadImplCaller> selfShared = MakeSharedRef(); + auto startCaller = [selfShared, deadline]() mutable { + selfShared->CallStart(deadline); + }; + Y_VERIFY(this->PQLib->GetQueuePool().GetQueue(this->QueueTag()).AddFunc(startCaller)); + } + + void CallStart(TInstant deadline) { + auto startPromise = StartPromise; + TImpl::Start(deadline).Subscribe([startPromise](const auto& future) mutable { + startPromise.SetValue(future.GetValue()); + }); + } + +private: + NThreading::TPromise<TStartResponse> StartPromise; + NThreading::TPromise<TError> DeadPromise = NThreading::NewPromise<TError>(); + bool IsDeadCalled = false; +}; + +#define LOCAL_CALLER(Func, Args, DeclTail, ...) \ + void Func Args DeclTail { \ + auto guard = Guard(this->DestroyLock); \ + if (this->IsCanceling) { \ + guard.Release(); \ + this->MakeCanceledError(__VA_ARGS__); \ + return; \ + } \ + auto selfShared = MakeSharedRef(); \ + auto caller = [selfShared, __VA_ARGS__]() mutable { \ + selfShared->TImpl::Func(__VA_ARGS__); \ + }; \ + Y_VERIFY(this->PQLib->GetQueuePool().GetQueue(this->QueueTag()).AddFunc(caller));\ + } + +// Assumes that there is function with promise interface +// first arg must be promise +#define LOCAL_CALLER_WITH_FUTURE(ReturnType, Func, Args, DeclTail, ...) \ + NThreading::TFuture<ReturnType> Func Args DeclTail { \ + auto promise = NThreading::NewPromise<ReturnType>(); \ + Func(__VA_ARGS__); \ + return promise.GetFuture(); \ + } + +// Helper classes for producer/consumer/processor impls. +// Wrappers that delegate calls to its methods to proper thread pool local threads. +// The result is that all calls to impl are serialized and locks are no more needed. +// +// Implies: +// That impl class is std::enable_shared_from_this descendant. + +template <class TImpl> +class TLocalProducerImplCaller: public TLocalStartDeadImplCaller<TImpl, TProducerCreateResponse> { +public: + using TLocalStartDeadImplCaller<TImpl, TProducerCreateResponse>::TLocalStartDeadImplCaller; + + LOCAL_CALLER(Write, + (NThreading::TPromise<TProducerCommitResponse>& promise, TProducerSeqNo seqNo, TData data), + noexcept override, + promise, seqNo, data) + + LOCAL_CALLER(Write, + (NThreading::TPromise<TProducerCommitResponse>& promise, TData data), + noexcept override, + promise, data) + + LOCAL_CALLER_WITH_FUTURE(TProducerCommitResponse, + Write, + (TProducerSeqNo seqNo, TData data), + noexcept override, + promise, seqNo, data) + + LOCAL_CALLER_WITH_FUTURE(TProducerCommitResponse, + Write, + (TData data), + noexcept override, + promise, data) + +protected: + using TLocalStartDeadImplCaller<TImpl, TProducerCreateResponse>::MakeCanceledError; + static void MakeCanceledError(NThreading::TPromise<TProducerCommitResponse>& promise, TProducerSeqNo seqNo, TData data) { + TWriteResponse responseProto; + MakeCanceledError(responseProto); + promise.SetValue(TProducerCommitResponse(seqNo, std::move(data), std::move(responseProto))); + } + + static void MakeCanceledError(NThreading::TPromise<TProducerCommitResponse>& promise, TData data) { + MakeCanceledError(promise, TProducerSeqNo(0), std::move(data)); + } + +private: + std::shared_ptr<TLocalProducerImplCaller> MakeSharedRef() { + return std::static_pointer_cast<TLocalProducerImplCaller>(this->shared_from_this()); + } +}; + +template <class TImpl> +class TLocalConsumerImplCaller: public TLocalStartDeadImplCaller<TImpl, TConsumerCreateResponse> { +public: + using TLocalStartDeadImplCaller<TImpl, TConsumerCreateResponse>::TLocalStartDeadImplCaller; + + LOCAL_CALLER(GetNextMessage, + (NThreading::TPromise<TConsumerMessage>& promise), + noexcept override, + promise) + + LOCAL_CALLER_WITH_FUTURE(TConsumerMessage, + GetNextMessage, + (), + noexcept override, + promise) + + LOCAL_CALLER(Commit, + (const TVector<ui64>& cookies), + noexcept override, + cookies) + + LOCAL_CALLER(RequestPartitionStatus, + (const TString& topic, ui64 partition, ui64 generation), + noexcept override, + topic, partition, generation) + +private: + using TLocalStartDeadImplCaller<TImpl, TConsumerCreateResponse>::MakeCanceledError; + static void MakeCanceledError(const TVector<ui64>& cookies) { + Y_UNUSED(cookies); + // Do nothing, because this method doesn't return future + } + + static void MakeCanceledError(const TString&, ui64, ui64) { + // Do nothing, because this method doesn't return future + } + + + std::shared_ptr<TLocalConsumerImplCaller> MakeSharedRef() { + return std::static_pointer_cast<TLocalConsumerImplCaller>(this->shared_from_this()); + } +}; + +template <class TImpl> +class TLocalProcessorImplCaller: public TCancelCaller<TImpl> { +public: + using TCancelCaller<TImpl>::TCancelCaller; + + LOCAL_CALLER(GetNextData, + (NThreading::TPromise<TOriginData>& promise), + noexcept override, + promise) + + LOCAL_CALLER_WITH_FUTURE(TOriginData, + GetNextData, + (), + noexcept override, + promise) + +private: + using TCancelCaller<TImpl>::MakeCanceledError; + static void MakeCanceledError(NThreading::TPromise<TOriginData>& promise) { + promise.SetValue(TOriginData()); // with empty data + } + + std::shared_ptr<TLocalProcessorImplCaller> MakeSharedRef() { + return std::static_pointer_cast<TLocalProcessorImplCaller>(this->shared_from_this()); + } +}; + +#undef LOCAL_CALLER +#undef LOCAL_CALLER_WITH_FUTURE + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/logger.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/logger.cpp new file mode 100644 index 0000000000..98d44d14c5 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/logger.cpp @@ -0,0 +1,40 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/logger.h> +#include <util/datetime/base.h> +#include <util/string/builder.h> +#include <util/stream/str.h> + +namespace NPersQueue { + +static const TStringBuf LogLevelsStrings[] = { + "EMERG", + "ALERT", + "CRITICAL_INFO", + "ERROR", + "WARNING", + "NOTICE", + "INFO", + "DEBUG", +}; + +TStringBuf TCerrLogger::LevelToString(int level) { + return LogLevelsStrings[ClampVal(level, 0, int(Y_ARRAY_SIZE(LogLevelsStrings) - 1))]; +} + +void TCerrLogger::Log(const TString& msg, const TString& sourceId, const TString& sessionId, int level) { + if (level > Level) { + return; + } + + TStringBuilder message; + message << TInstant::Now() << " :" << LevelToString(level) << ":"; + if (sourceId) { + message << " SourceId [" << sourceId << "]:"; + } + if (sessionId) { + message << " SessionId [" << sessionId << "]:"; + } + message << " " << msg << "\n"; + Cerr << message; +} + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_consumer.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_consumer.cpp new file mode 100644 index 0000000000..cd986532dc --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_consumer.cpp @@ -0,0 +1,645 @@ +#include "multicluster_consumer.h" +#include "persqueue_p.h" + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> + +#include <library/cpp/containers/intrusive_rb_tree/rb_tree.h> + +#include <util/generic/guid.h> +#include <util/generic/strbuf.h> + +namespace NPersQueue { + +struct TMultiClusterConsumer::TSubconsumerInfo { + TSubconsumerInfo(std::shared_ptr<IConsumerImpl> consumer); + + std::shared_ptr<IConsumerImpl> Consumer; + NThreading::TFuture<TConsumerCreateResponse> StartFuture; + NThreading::TFuture<TError> DeadFuture; + bool StartResponseWasGot = false; // We call methods of this consumer only after this flag was set. + + std::deque<NThreading::TFuture<TConsumerMessage>> PendingRequests; +}; + +namespace { + +struct TConsumerCookieKey; +struct TUserCookieMappingItem; +struct TConsumerCookieMappingItem; + +struct TConsumerCookieKey { + size_t Subconsumer = 0; + ui64 ConsumerCookie = 0; + + TConsumerCookieKey(size_t subconsumer, ui64 consumerCookie) + : Subconsumer(subconsumer) + , ConsumerCookie(consumerCookie) + { + } + + bool operator<(const TConsumerCookieKey& other) const { + return std::make_pair(Subconsumer, ConsumerCookie) < std::make_pair(other.Subconsumer, other.ConsumerCookie); + } +}; + +struct TUserCookieCmp { + static ui64 GetKey(const TUserCookieMappingItem& item); + + static bool Compare(ui64 l, const TUserCookieMappingItem& r) { + return l < GetKey(r); + } + + static bool Compare(const TUserCookieMappingItem& l, ui64 r) { + return GetKey(l) < r; + } + + static bool Compare(const TUserCookieMappingItem& l, const TUserCookieMappingItem& r) { + return GetKey(l) < GetKey(r); + } +}; + +struct TConsumerCookieCmp { + static TConsumerCookieKey GetKey(const TConsumerCookieMappingItem& item); + + static bool Compare(const TConsumerCookieKey& l, const TConsumerCookieMappingItem& r) { + return l < GetKey(r); + } + + static bool Compare(const TConsumerCookieMappingItem& l, const TConsumerCookieKey& r) { + return GetKey(l) < r; + } + + static bool Compare(const TConsumerCookieMappingItem& l, const TConsumerCookieMappingItem& r) { + return GetKey(l) < GetKey(r); + } +}; + +struct TUserCookieMappingItem : public TRbTreeItem<TUserCookieMappingItem, TUserCookieCmp> { +}; + +struct TConsumerCookieMappingItem : public TRbTreeItem<TConsumerCookieMappingItem, TConsumerCookieCmp> { +}; + +bool IsRetryable(const grpc::Status& status) { + switch (status.error_code()) { + case grpc::OK: + case grpc::CANCELLED: + case grpc::INVALID_ARGUMENT: + case grpc::NOT_FOUND: + case grpc::ALREADY_EXISTS: + case grpc::PERMISSION_DENIED: + case grpc::UNAUTHENTICATED: + case grpc::FAILED_PRECONDITION: + case grpc::ABORTED: + case grpc::OUT_OF_RANGE: + case grpc::UNIMPLEMENTED: + return false; + + case grpc::UNKNOWN: + case grpc::DEADLINE_EXCEEDED: + case grpc::RESOURCE_EXHAUSTED: + case grpc::INTERNAL: + case grpc::UNAVAILABLE: + case grpc::DATA_LOSS: + case grpc::DO_NOT_USE: + return true; + } +} + +} // namespace + +struct TMultiClusterConsumer::TCookieMappingItem : public TUserCookieMappingItem, public TConsumerCookieMappingItem { + TCookieMappingItem(size_t subconsumer, ui64 consumerCookie, ui64 userCookie) + : Subconsumer(subconsumer) + , ConsumerCookie(consumerCookie) + , UserCookie(userCookie) + { + } + + size_t Subconsumer = 0; + ui64 ConsumerCookie = 0; + ui64 UserCookie = 0; +}; + +namespace { + +ui64 TUserCookieCmp::GetKey(const TUserCookieMappingItem& item) { + return static_cast<const TMultiClusterConsumer::TCookieMappingItem&>(item).UserCookie; +} + +TConsumerCookieKey TConsumerCookieCmp::GetKey(const TConsumerCookieMappingItem& item) { + const auto& src = static_cast<const TMultiClusterConsumer::TCookieMappingItem&>(item); + return { src.Subconsumer, src.ConsumerCookie }; +} + +} // namespace + +struct TMultiClusterConsumer::TCookieMapping { + using TUserCookieTree = TRbTree<TUserCookieMappingItem, TUserCookieCmp>; + using TConsumerCookieTree = TRbTree<TConsumerCookieMappingItem, TConsumerCookieCmp>; + + struct TConsumerCookieDestroy : public TConsumerCookieTree::TDestroy { + void operator()(TConsumerCookieMappingItem& item) const { + TDestroy::operator()(item); // Remove from tree + delete static_cast<TCookieMappingItem*>(&item); + } + }; + + ~TCookieMapping() { + UserCookieTree.ForEachNoOrder(TUserCookieTree::TDestroy()); + UserCookieTree.Init(); + + ConsumerCookieTree.ForEachNoOrder(TConsumerCookieDestroy()); + ConsumerCookieTree.Init(); + } + + ui64 AddMapping(size_t subconsumer, ui64 cookie) { // Returns user cookie + auto* newItem = new TCookieMappingItem(subconsumer, cookie, NextCookie++); + UserCookieTree.Insert(newItem); + ConsumerCookieTree.Insert(newItem); + return newItem->UserCookie; + } + + TCookieMappingItem* FindMapping(ui64 userCookie) { + return DownCast(UserCookieTree.Find(userCookie)); + } + + TCookieMappingItem* FindMapping(size_t subconsumer, ui64 cookie) { + return DownCast(ConsumerCookieTree.Find(TConsumerCookieKey(subconsumer, cookie))); + } + + void RemoveUserCookieMapping(TCookieMappingItem* item) { + Y_ASSERT(item); + static_cast<TUserCookieMappingItem*>(item)->UnLink(); + } + + void RemoveMapping(TCookieMappingItem* item) { + Y_ASSERT(item); + Y_ASSERT(!static_cast<TUserCookieMappingItem*>(item)->ParentTree()); + static_cast<TUserCookieMappingItem*>(item)->UnLink(); // Just in case + static_cast<TConsumerCookieMappingItem*>(item)->UnLink(); + delete item; + } + +private: + template <class T> + static TCookieMappingItem* DownCast(T* item) { + if (item) { + return static_cast<TCookieMappingItem*>(item); + } else { + return nullptr; + } + } + + TUserCookieTree UserCookieTree; + TConsumerCookieTree ConsumerCookieTree; + ui64 NextCookie = 1; +}; + +TMultiClusterConsumer::TMultiClusterConsumer(const TConsumerSettings& settings, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger) + : IConsumerImpl(std::move(destroyEventRef), std::move(pqLib)) + , Settings(settings) + , Logger(std::move(logger)) + , SessionId(CreateGuidAsString()) + , CookieMapping(MakeHolder<TCookieMapping>()) +{ + PatchSettings(); +} + +TMultiClusterConsumer::~TMultiClusterConsumer() { + Destroy(); +} + +void TMultiClusterConsumer::PatchSettings() { + if (!Settings.ReconnectOnFailure) { + WARN_LOG("Ignoring ReconnectOnFailure=false option for multicluster consumer", "", SessionId); + Settings.ReconnectOnFailure = true; + } + if (Settings.MaxAttempts != std::numeric_limits<unsigned>::max()) { + WARN_LOG("Ignoring MaxAttempts option for multicluster consumer", "", SessionId); + Settings.MaxAttempts = std::numeric_limits<unsigned>::max(); + } +} + +NThreading::TFuture<TConsumerCreateResponse> TMultiClusterConsumer::Start(TInstant deadline) noexcept { + Y_VERIFY(State == EState::Created); + StartClusterDiscovery(deadline); + return StartPromise.GetFuture(); +} + +NThreading::TFuture<TError> TMultiClusterConsumer::IsDead() noexcept { + return IsDeadPromise.GetFuture(); +} + +void TMultiClusterConsumer::ScheduleClusterDiscoveryRetry(TInstant deadline) { + const TInstant now = TInstant::Now(); + if (ClusterDiscoveryAttemptsDone >= Settings.MaxAttempts || now >= deadline) { + Destroy(TStringBuilder() << "Failed " << ClusterDiscoveryAttemptsDone << " cluster discovery attempts", NErrorCode::CREATE_TIMEOUT); + return; + } + + if (!IsRetryable(ClusterDiscoverResult->first)) { + Destroy(TStringBuilder() << "Cluster discovery failed: " << static_cast<int>(ClusterDiscoverResult->first.error_code())); + return; + } + + TDuration delay = Min(Settings.MaxReconnectionDelay, ClusterDiscoveryAttemptsDone * Settings.ReconnectionDelay, deadline - now); + std::weak_ptr<TMultiClusterConsumer> weakRef = shared_from_this(); + PQLib->GetScheduler().Schedule(delay, this, [weakRef, deadline]() { + if (auto self = weakRef.lock()) { + self->StartClusterDiscovery(deadline); + } + }); +} + +void TMultiClusterConsumer::StartClusterDiscovery(TInstant deadline) { + if (State == EState::Dead) { + return; + } + + State = EState::WaitingClusterDiscovery; + ++ClusterDiscoveryAttemptsDone; + DEBUG_LOG("Starting cluster discovery", "", SessionId); + + auto discoverer = MakeIntrusive<TConsumerChannelOverCdsImpl>(Settings, PQLib.Get(), Logger); + discoverer->Start(); + ClusterDiscoverResult = discoverer->GetResultPtr(); + ClusterDiscoverer = std::move(discoverer); + + std::weak_ptr<TMultiClusterConsumer> self = shared_from_this(); + auto handler = [self, deadline](const auto&) { + if (auto selfShared = self.lock()) { + selfShared->OnClusterDiscoveryDone(deadline); + } + }; + PQLib->Subscribe(ClusterDiscoverer->GetChannel(), this, handler); +} + +void TMultiClusterConsumer::OnClusterDiscoveryDone(TInstant deadline) { + if (State != EState::WaitingClusterDiscovery) { + return; + } + + ClusterDiscoverer = nullptr; // delete + if (!ClusterDiscoverResult->first.ok()) { + INFO_LOG("Failed to discover clusters. Grpc error: " << static_cast<int>(ClusterDiscoverResult->first.error_code()), "", SessionId); + ScheduleClusterDiscoveryRetry(deadline); // Destroys if we shouldn't retry. + return; + } + + if (static_cast<size_t>(ClusterDiscoverResult->second.read_sessions_clusters_size()) != Settings.Topics.size()) { + Destroy("Got unexpected cluster discovery result"); + return; + } + + if (deadline != TInstant::Max()) { + const TInstant now = TInstant::Now(); + if (now > deadline) { + Destroy("Start timeout", NErrorCode::CREATE_TIMEOUT); + return; + } + + std::weak_ptr<TMultiClusterConsumer> self = shared_from_this(); + auto onStartTimeout = [self] { + if (auto selfShared = self.lock()) { + selfShared->OnStartTimeout(); + } + }; + StartDeadlineCallback = + PQLib->GetScheduler().Schedule(deadline, this, onStartTimeout); + } + + State = EState::StartingSubconsumers; + DEBUG_LOG("Starting subconsumers", "", SessionId); + + // Group topics by clusters. + THashMap<TString, TVector<TString>> clusterToTopics; + for (size_t topicIndex = 0, topicsCount = Settings.Topics.size(); topicIndex < topicsCount; ++topicIndex) { + const Ydb::PersQueue::ClusterDiscovery::ReadSessionClusters& clusters = ClusterDiscoverResult->second.read_sessions_clusters(topicIndex); + const TString& topic = Settings.Topics[topicIndex]; + for (const Ydb::PersQueue::ClusterDiscovery::ClusterInfo& clusterInfo : clusters.clusters()) { + clusterToTopics[clusterInfo.endpoint()].push_back(topic); + } + } + + if (clusterToTopics.empty()) { + Destroy("Got empty endpoint set from cluster discovery"); + return; + } + + // Start consumer on every cluster with its topics set. + for (auto&& [clusterEndpoint, topics] : clusterToTopics) { + // Settings. + TConsumerSettings settings = Settings; + settings.Server = ApplyClusterEndpoint(settings.Server, clusterEndpoint); + settings.ReadFromAllClusterSources = false; + settings.Topics = std::move(topics); + settings.Unpack = false; // Unpack is being done in upper level decompressing consumer. + settings.MaxMemoryUsage = Max(settings.MaxMemoryUsage / clusterToTopics.size(), static_cast<size_t>(1)); + if (settings.MaxUncommittedSize > 0) { // Limit is enabled. + settings.MaxUncommittedSize = Max(settings.MaxUncommittedSize / clusterToTopics.size(), static_cast<size_t>(1)); + } + if (settings.MaxUncommittedCount > 0) { // Limit is enabled. + settings.MaxUncommittedCount = Max(settings.MaxUncommittedCount / clusterToTopics.size(), static_cast<size_t>(1)); + } + const size_t subconsumerIndex = Subconsumers.size(); + + // Create subconsumer. + Subconsumers.push_back(PQLib->CreateRawRetryingConsumer(settings, DestroyEventRef, Logger)); + + // Subscribe on start. + std::weak_ptr<TMultiClusterConsumer> self = shared_from_this(); + auto handler = [self, subconsumerIndex](const auto&) { + if (auto selfShared = self.lock()) { + selfShared->OnSubconsumerStarted(subconsumerIndex); + } + }; + PQLib->Subscribe(Subconsumers.back().StartFuture, this, handler); + } +} + +void TMultiClusterConsumer::OnStartTimeout() { + if (State == EState::Dead || State == EState::Working) { + return; + } + + StartDeadlineCallback = nullptr; + Destroy("Start timeout", NErrorCode::CREATE_TIMEOUT); +} + + +void TMultiClusterConsumer::OnSubconsumerStarted(size_t subconsumerIndex) { + if (State == EState::Dead) { + return; + } + + auto& subconsumerInfo = Subconsumers[subconsumerIndex]; + const auto& result = subconsumerInfo.StartFuture.GetValue(); + if (result.Response.HasError()) { + WARN_LOG("Got error on starting subconsumer: " << result.Response.GetError(), "", SessionId); + Destroy(result.Response.GetError()); + return; + } + + // Process start response + subconsumerInfo.StartResponseWasGot = true; + INFO_LOG("Subconsumer " << subconsumerIndex << " started session with id " << result.Response.GetInit().GetSessionId(), "", SessionId); + + // Move to new state + if (State == EState::StartingSubconsumers) { + State = EState::Working; + if (StartDeadlineCallback) { + StartDeadlineCallback->TryCancel(); + StartDeadlineCallback = nullptr; + } + + { + TReadResponse resp; + resp.MutableInit()->SetSessionId(SessionId); + StartPromise.SetValue(TConsumerCreateResponse(std::move(resp))); + } + } else { + RequestSubconsumers(); // Make requests for new subconsumer. + } +} + +void TMultiClusterConsumer::Commit(const TVector<ui64>& cookies) noexcept { + if (State != EState::Working) { + Destroy("Requesting commit, but consumer is not in working state", NErrorCode::BAD_REQUEST); + return; + } + + TVector<TVector<ui64>> subconsumersCookies(Subconsumers.size()); + for (ui64 userCookie : cookies) { + TCookieMappingItem* mapping = CookieMapping->FindMapping(userCookie); + if (!mapping) { + Destroy(TStringBuilder() << "Wrong cookie " << userCookie, NErrorCode::WRONG_COOKIE); + return; + } + Y_ASSERT(mapping->Subconsumer < Subconsumers.size()); + subconsumersCookies[mapping->Subconsumer].push_back(mapping->ConsumerCookie); + CookieMapping->RemoveUserCookieMapping(mapping); // Avoid double commit. This will ensure error on second commit with the same cookie. + } + + for (size_t subconsumerIndex = 0; subconsumerIndex < Subconsumers.size(); ++subconsumerIndex) { + const TVector<ui64>& subconsumerCommitRequest = subconsumersCookies[subconsumerIndex]; + if (!subconsumerCommitRequest.empty()) { + Subconsumers[subconsumerIndex].Consumer->Commit(subconsumerCommitRequest); + } + } +} + +void TMultiClusterConsumer::GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept { + if (State != EState::Working) { + promise.SetValue(MakeResponse(MakeError("Requesting next message, but consumer is not in working state", NErrorCode::BAD_REQUEST))); + return; + } + + Requests.push_back(promise); + CheckReadyResponses(); + RequestSubconsumers(); +} + +void TMultiClusterConsumer::RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept { + if (State != EState::Working) { + Destroy("Requesting partition status, but consumer is not in working state", NErrorCode::BAD_REQUEST); + return; + } + + const auto subconsumerIdIt = OldTopicName2Subconsumer.find(topic); + if (subconsumerIdIt != OldTopicName2Subconsumer.end()) { + Y_ASSERT(subconsumerIdIt->second < Subconsumers.size()); + Subconsumers[subconsumerIdIt->second].Consumer->RequestPartitionStatus(topic, partition, generation); + } else { + WARN_LOG("Requested partition status for topic \"" << topic << "\" (partition " << partition << ", generation " << generation << "), but there is no such lock session. Ignoring request", "", SessionId); + } +} + +void TMultiClusterConsumer::CheckReadyResponses() { + const size_t prevCurrentSubconsumer = CurrentSubconsumer; + + do { + if (!ProcessReadyResponses(CurrentSubconsumer)) { + break; + } + + // Next subconsumer in round robin way. + ++CurrentSubconsumer; + if (CurrentSubconsumer == Subconsumers.size()) { + CurrentSubconsumer = 0; + } + } while (CurrentSubconsumer != prevCurrentSubconsumer && !Requests.empty()); +} + +bool TMultiClusterConsumer::ProcessReadyResponses(size_t subconsumerIndex) { + if (Subconsumers.empty()) { + Y_VERIFY(State == EState::Dead); + return false; + } + Y_VERIFY(subconsumerIndex < Subconsumers.size()); + TSubconsumerInfo& consumerInfo = Subconsumers[subconsumerIndex]; + while (!consumerInfo.PendingRequests.empty() && !Requests.empty() && consumerInfo.PendingRequests.front().HasValue()) { + if (!TranslateConsumerMessage(Requests.front(), consumerInfo.PendingRequests.front().ExtractValue(), subconsumerIndex)) { + return false; + } + Requests.pop_front(); + consumerInfo.PendingRequests.pop_front(); + } + return true; +} + +bool TMultiClusterConsumer::TranslateConsumerMessage(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex) { + switch (subconsumerResponse.Type) { + case EMT_LOCK: + return TranslateConsumerMessageLock(promise, std::move(subconsumerResponse), subconsumerIndex); + case EMT_RELEASE: + return TranslateConsumerMessageRelease(promise, std::move(subconsumerResponse), subconsumerIndex); + case EMT_DATA: + return TranslateConsumerMessageData(promise, std::move(subconsumerResponse), subconsumerIndex); + case EMT_ERROR: + return TranslateConsumerMessageError(promise, std::move(subconsumerResponse), subconsumerIndex); + case EMT_STATUS: + return TranslateConsumerMessageStatus(promise, std::move(subconsumerResponse), subconsumerIndex); + case EMT_COMMIT: + return TranslateConsumerMessageCommit(promise, std::move(subconsumerResponse), subconsumerIndex); + } +} + +bool TMultiClusterConsumer::TranslateConsumerMessageData(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex) { + auto* data = subconsumerResponse.Response.MutableData(); + data->SetCookie(CookieMapping->AddMapping(subconsumerIndex, data->GetCookie())); + promise.SetValue(std::move(subconsumerResponse)); + return true; +} + +bool TMultiClusterConsumer::TranslateConsumerMessageCommit(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex) { + for (ui64& consumerCookie : *subconsumerResponse.Response.MutableCommit()->MutableCookie()) { + auto* mapping = CookieMapping->FindMapping(subconsumerIndex, consumerCookie); + if (!mapping) { + Destroy(TStringBuilder() << "Received unknown cookie " << consumerCookie << " commit"); + return false; + } + Y_VERIFY(mapping); + consumerCookie = mapping->UserCookie; + CookieMapping->RemoveMapping(mapping); // Invalidates mapping + } + promise.SetValue(std::move(subconsumerResponse)); + return true; +} + +bool TMultiClusterConsumer::TranslateConsumerMessageLock(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex) { + OldTopicName2Subconsumer[subconsumerResponse.Response.GetLock().GetTopic()] = subconsumerIndex; + promise.SetValue(std::move(subconsumerResponse)); + return true; +} + +bool TMultiClusterConsumer::TranslateConsumerMessageRelease(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex) { + Y_UNUSED(subconsumerIndex); + promise.SetValue(std::move(subconsumerResponse)); + return true; +} + +bool TMultiClusterConsumer::TranslateConsumerMessageError(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex) { + WARN_LOG("Got error from subconsumer " << subconsumerIndex << ": " << subconsumerResponse.Response.GetError(), "", SessionId); + Destroy(subconsumerResponse.Response.GetError()); + Y_UNUSED(promise); + return false; +} + +bool TMultiClusterConsumer::TranslateConsumerMessageStatus(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex) { + Y_UNUSED(subconsumerIndex); + promise.SetValue(std::move(subconsumerResponse)); + return true; +} + +void TMultiClusterConsumer::RequestSubconsumers() { + const size_t inflight = Requests.size(); + if (!inflight) { + return; + } + TMaybe<std::weak_ptr<TMultiClusterConsumer>> maybeSelf; + for (size_t subconsumerIndex = 0; subconsumerIndex < Subconsumers.size(); ++subconsumerIndex) { + TSubconsumerInfo& consumerInfo = Subconsumers[subconsumerIndex]; + if (!consumerInfo.StartResponseWasGot) { + continue; // Consumer hasn't started yet. + } + while (consumerInfo.PendingRequests.size() < inflight) { + if (!maybeSelf) { + maybeSelf.ConstructInPlace(shared_from_this()); // Don't construct it if we don't need it. + } + consumerInfo.PendingRequests.push_back(consumerInfo.Consumer->GetNextMessage()); + PQLib->Subscribe(consumerInfo.PendingRequests.back(), + this, + [self = *maybeSelf, subconsumerIndex](const auto&) { + if (auto selfShared = self.lock()) { + selfShared->ProcessReadyResponses(subconsumerIndex); + } + }); + } + } +} + +void TMultiClusterConsumer::Destroy(const TError& description) { + if (State == EState::Dead) { + return; + } + State = EState::Dead; + + WARN_LOG("Destroying consumer with error description: " << description, "", SessionId); + + StartPromise.TrySetValue(MakeCreateResponse(description)); + + if (StartDeadlineCallback) { + StartDeadlineCallback->TryCancel(); + } + + for (auto& reqPromise : Requests) { + reqPromise.SetValue(MakeResponse(description)); + } + Requests.clear(); + Subconsumers.clear(); + + IsDeadPromise.SetValue(description); + + DestroyPQLibRef(); +} + +void TMultiClusterConsumer::Destroy(const TString& description, NErrorCode::EErrorCode code) { + Destroy(MakeError(description, code)); +} + +void TMultiClusterConsumer::Destroy() { + Destroy(GetCancelReason()); +} + +void TMultiClusterConsumer::Cancel() { + Destroy(GetCancelReason()); +} + +TError TMultiClusterConsumer::MakeError(const TString& description, NErrorCode::EErrorCode code) { + TError error; + error.SetDescription(description); + error.SetCode(code); + return error; +} + +TConsumerCreateResponse TMultiClusterConsumer::MakeCreateResponse(const TError& description) { + TReadResponse res; + res.MutableError()->CopyFrom(description); + return TConsumerCreateResponse(std::move(res)); +} + +TConsumerMessage TMultiClusterConsumer::MakeResponse(const TError& description) { + TReadResponse res; + res.MutableError()->CopyFrom(description); + return TConsumerMessage(std::move(res)); +} + +TMultiClusterConsumer::TSubconsumerInfo::TSubconsumerInfo(std::shared_ptr<IConsumerImpl> consumer) + : Consumer(std::move(consumer)) + , StartFuture(Consumer->Start()) + , DeadFuture(Consumer->IsDead()) +{ +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_consumer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_consumer.h new file mode 100644 index 0000000000..ab8a338fbe --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_consumer.h @@ -0,0 +1,93 @@ +#pragma once +#include "channel_p.h" +#include "scheduler.h" +#include "iconsumer_p.h" +#include "internals.h" + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> + +namespace NPersQueue { + +class TPQLibPrivate; + +class TMultiClusterConsumer: public IConsumerImpl, public std::enable_shared_from_this<TMultiClusterConsumer> { +public: + struct TSubconsumerInfo; + struct TCookieMappingItem; + struct TCookieMapping; + +public: + TMultiClusterConsumer(const TConsumerSettings& settings, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger); + ~TMultiClusterConsumer(); + + NThreading::TFuture<TConsumerCreateResponse> Start(TInstant deadline) noexcept override; + NThreading::TFuture<TError> IsDead() noexcept override; + + void GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept override; + void Commit(const TVector<ui64>& cookies) noexcept override; + void RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept override; + + void Cancel() override; + +private: + void PatchSettings(); + void StartClusterDiscovery(TInstant deadline); + void OnClusterDiscoveryDone(TInstant deadline); + void OnStartTimeout(); + void OnSubconsumerStarted(size_t subconsumerIndex); + + void CheckReadyResponses(); + bool ProcessReadyResponses(size_t subconsumerIndex); // False if consumer has been destroyed. + void RequestSubconsumers(); // Make requests inflight in every subconsumer equal to our superconsumer inflight. + + bool TranslateConsumerMessage(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex); // False if consumer has been destroyed. + bool TranslateConsumerMessageLock(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex); + bool TranslateConsumerMessageRelease(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex); + bool TranslateConsumerMessageData(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex); + bool TranslateConsumerMessageCommit(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex); + bool TranslateConsumerMessageError(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex); // Calls destroy on consumer + bool TranslateConsumerMessageStatus(NThreading::TPromise<TConsumerMessage>& promise, TConsumerMessage&& subconsumerResponse, size_t subconsumerIndex); + + void Destroy(const TError& description); + void Destroy(const TString& description, NErrorCode::EErrorCode code = NErrorCode::ERROR); + void Destroy(); + + static TError MakeError(const TString& description, NErrorCode::EErrorCode code = NErrorCode::ERROR); + static TConsumerCreateResponse MakeCreateResponse(const TError& description); + static TConsumerMessage MakeResponse(const TError& description); + + void ScheduleClusterDiscoveryRetry(TInstant deadline); + +private: + TConsumerSettings Settings; + TIntrusivePtr<ILogger> Logger; + TString SessionId; + + NThreading::TPromise<TConsumerCreateResponse> StartPromise = NThreading::NewPromise<TConsumerCreateResponse>(); + NThreading::TPromise<TError> IsDeadPromise = NThreading::NewPromise<TError>(); + + enum class EState { + Created, // Before Start() was created. + WaitingClusterDiscovery, // After Start() was called and before we received CDS response. + StartingSubconsumers, // After we received CDS response and before consumer actually started. + Working, // When one or more consumers had started. + Dead, + }; + EState State = EState::Created; + + TIntrusivePtr<TScheduler::TCallbackHandler> StartDeadlineCallback; + TChannelImplPtr ClusterDiscoverer; + TConsumerChannelOverCdsImpl::TResultPtr ClusterDiscoverResult; + size_t ClusterDiscoveryAttemptsDone = 0; + + std::vector<TSubconsumerInfo> Subconsumers; + size_t CurrentSubconsumer = 0; // Index of current subconsumer to take responses from + + THolder<TCookieMapping> CookieMapping; + THashMap<TString, size_t> OldTopicName2Subconsumer; // Used in lock session for requesting partition status. + + std::deque<NThreading::TPromise<TConsumerMessage>> Requests; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_consumer_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_consumer_ut.cpp new file mode 100644 index 0000000000..a29a18cbaf --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_consumer_ut.cpp @@ -0,0 +1,227 @@ +#include "multicluster_consumer.h" +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/sdk_test_setup.h> + +#include <library/cpp/testing/unittest/registar.h> + +#include <util/system/event.h> + +namespace NPersQueue { + +Y_UNIT_TEST_SUITE(TMultiClusterConsumerTest) { + TConsumerSettings MakeFakeConsumerSettings() { + TConsumerSettings settings; + settings.Topics.push_back("topic"); + settings.ReadFromAllClusterSources = true; + settings.ReadMirroredPartitions = false; + settings.MaxAttempts = 1; // Consumer should fix this setting. Check it. + return settings; + } + + TConsumerSettings MakeConsumerSettings(const SDKTestSetup& setup) { + TConsumerSettings settings = setup.GetConsumerSettings(); + settings.ReadFromAllClusterSources = true; + settings.ReadMirroredPartitions = false; + settings.MaxAttempts = 1; // Consumer should fix this setting. Check it. + return settings; + } + + Y_UNIT_TEST(NotStartedConsumerCanBeDestructed) { + // Test that consumer doesn't hang on till shutdown + TPQLib lib; + lib.CreateConsumer(MakeFakeConsumerSettings(), new TCerrLogger(TLOG_DEBUG)); + } + + Y_UNIT_TEST(StartedConsumerCanBeDestructed) { + // Test that consumer doesn't hang on till shutdown + TPQLib lib; + auto consumer = lib.CreateConsumer(MakeFakeConsumerSettings(), new TCerrLogger(TLOG_DEBUG)); + auto isDead = consumer->IsDead(); + consumer->Start(); + consumer = nullptr; + isDead.GetValueSync(); + } + + Y_UNIT_TEST(CommitRightAfterStart) { + SDKTestSetup setup("InfiniteStart"); + TPQLib lib; + auto settings = MakeConsumerSettings(setup); + settings.Topics[0] = "unknown_topic"; + auto consumer = lib.CreateConsumer(settings, new TCerrLogger(TLOG_DEBUG)); + consumer->Start(); + auto isDead = consumer->IsDead(); + consumer->Commit({42}); + if (GrpcV1EnabledByDefault()) { + UNIT_ASSERT_EQUAL_C(isDead.GetValueSync().GetCode(), NErrorCode::WRONG_COOKIE, isDead.GetValueSync()); + } else { + UNIT_ASSERT_STRINGS_EQUAL_C(isDead.GetValueSync().GetDescription(), "Requesting commit, but consumer is not in working state", isDead.GetValueSync()); + } + } + + Y_UNIT_TEST(GetNextMessageRightAfterStart) { + if (GrpcV1EnabledByDefault()) { + return; + } + SDKTestSetup setup("InfiniteStart"); + TPQLib lib; + auto settings = MakeConsumerSettings(setup); + settings.Topics[0] = "unknown_topic"; + auto consumer = lib.CreateConsumer(settings, new TCerrLogger(TLOG_DEBUG)); + consumer->Start(); + auto nextMessage = consumer->GetNextMessage(); + UNIT_ASSERT_STRINGS_EQUAL(nextMessage.GetValueSync().Response.GetError().GetDescription(), "Requesting next message, but consumer is not in working state"); + } + + Y_UNIT_TEST(RequestPartitionStatusRightAfterStart) { + if (GrpcV1EnabledByDefault()) { + return; + } + SDKTestSetup setup("InfiniteStart"); + TPQLib lib; + auto settings = MakeConsumerSettings(setup); + settings.Topics[0] = "unknown_topic"; + auto consumer = lib.CreateConsumer(settings, new TCerrLogger(TLOG_DEBUG)); + consumer->Start(); + auto isDead = consumer->IsDead(); + consumer->RequestPartitionStatus("topic", 42, 42); + UNIT_ASSERT_STRINGS_EQUAL(isDead.GetValueSync().GetDescription(), "Requesting partition status, but consumer is not in working state"); + } + + Y_UNIT_TEST(StartedConsumerCantBeStartedAgain) { + TPQLib lib; + auto consumer = lib.CreateConsumer(MakeFakeConsumerSettings(), new TCerrLogger(TLOG_DEBUG)); + auto startFuture = consumer->Start(); + auto startFuture2 = consumer->Start(); + UNIT_ASSERT(startFuture2.HasValue()); + UNIT_ASSERT_STRINGS_EQUAL(startFuture2.GetValueSync().Response.GetError().GetDescription(), "Start was already called"); + } + + Y_UNIT_TEST(StartWithFailedCds) { + TPQLib lib; + auto consumer = lib.CreateConsumer(MakeFakeConsumerSettings(), new TCerrLogger(TLOG_DEBUG)); + auto startFuture = consumer->Start(TDuration::MilliSeconds(10)); + if (!GrpcV1EnabledByDefault()) { + UNIT_ASSERT_STRING_CONTAINS(startFuture.GetValueSync().Response.GetError().GetDescription(), " cluster discovery attempts"); + } + auto isDead = consumer->IsDead(); + isDead.Wait(); + } + + Y_UNIT_TEST(StartDeadline) { + if (GrpcV1EnabledByDefault()) { + return; + } + TPQLib lib; + TConsumerSettings settings = MakeFakeConsumerSettings(); + auto consumer = lib.CreateConsumer(settings, new TCerrLogger(TLOG_DEBUG)); + auto startFuture = consumer->Start(TDuration::MilliSeconds(100)); + UNIT_ASSERT_EQUAL_C(startFuture.GetValueSync().Response.GetError().GetCode(), NErrorCode::CREATE_TIMEOUT, startFuture.GetValueSync().Response); + } + + Y_UNIT_TEST(StartDeadlineAfterCds) { + SDKTestSetup setup("StartDeadlineAfterCds"); + TPQLib lib; + auto settings = MakeConsumerSettings(setup); + settings.Topics[0] = "unknown_topic"; + auto consumer = lib.CreateConsumer(settings, new TCerrLogger(TLOG_DEBUG)); + auto startFuture = consumer->Start(TDuration::MilliSeconds(1)); + if (!GrpcV1EnabledByDefault()) { + UNIT_ASSERT_STRINGS_EQUAL(startFuture.GetValueSync().Response.GetError().GetDescription(), "Start timeout"); + } + auto isDead = consumer->IsDead(); + isDead.Wait(); + } + + void WorksWithSingleDc(bool oneDcIsUnavailable) { + SDKTestSetup setup("WorksWithSingleDc", false); + if (!oneDcIsUnavailable) { + setup.SetSingleDataCenter(); // By default one Dc is fake in test setup. This call overrides it. + } + setup.Start(); + auto consumer = setup.StartConsumer(MakeConsumerSettings(setup)); + setup.WriteToTopic({"msg1", "msg2"}); + setup.ReadFromTopic({{"msg1", "msg2"}}, true, consumer.Get()); + } + + Y_UNIT_TEST(WorksWithSingleDc) { + WorksWithSingleDc(false); + } + + Y_UNIT_TEST(WorksWithSingleDcWhileSecondDcIsUnavailable) { + WorksWithSingleDc(true); + } + + Y_UNIT_TEST(FailsOnCommittingWrongCookie) { + SDKTestSetup setup("FailsOnCommittingWrongCookie"); + auto consumer = setup.StartConsumer(MakeConsumerSettings(setup)); + setup.WriteToTopic({"msg1", "msg2"}); + auto nextMessage = consumer->GetNextMessage(); + UNIT_ASSERT_EQUAL(nextMessage.GetValueSync().Type, EMT_DATA); + auto isDead = consumer->IsDead(); + consumer->Commit({nextMessage.GetValueSync().Response.GetData().GetCookie() + 1}); + UNIT_ASSERT_STRING_CONTAINS(isDead.GetValueSync().GetDescription(), "Wrong cookie"); + } + + // Test that consumer is properly connected with decompressing consumer. + void Unpacks(bool unpack) { + SDKTestSetup setup(TStringBuilder() << "Unpacks(" << (unpack ? "true)" : "false)")); + setup.WriteToTopic({"msg"}); + auto settings = MakeConsumerSettings(setup); + settings.Unpack = unpack; + auto consumer = setup.StartConsumer(settings); + auto readResponse = consumer->GetNextMessage().GetValueSync(); + UNIT_ASSERT_C(!readResponse.Response.HasError(), readResponse.Response); + UNIT_ASSERT_C(readResponse.Response.HasData(), readResponse.Response); + if (unpack) { + UNIT_ASSERT_EQUAL_C(readResponse.Response.GetData().GetMessageBatch(0).GetMessage(0).GetMeta().GetCodec(), NPersQueueCommon::RAW, readResponse.Response); + UNIT_ASSERT_STRINGS_EQUAL(readResponse.Response.GetData().GetMessageBatch(0).GetMessage(0).GetData(), "msg"); + } else { + UNIT_ASSERT_EQUAL_C(readResponse.Response.GetData().GetMessageBatch(0).GetMessage(0).GetMeta().GetCodec(), NPersQueueCommon::GZIP, readResponse.Response); + } + } + + Y_UNIT_TEST(UnpacksIfNeeded) { + Unpacks(true); + } + + Y_UNIT_TEST(DoesNotUnpackIfNeeded) { + Unpacks(false); + } + + Y_UNIT_TEST(WorksWithManyDc) { + SDKTestSetup setup1("WorksWithManyDc1", false); + SDKTestSetup setup2("WorksWithManyDc2"); + SDKTestSetup setup3("WorksWithManyDc3", false); + setup1.AddDataCenter("dc2", setup2); + setup1.AddDataCenter("dc3", setup3); + setup3.GetGrpcServerOptions().SetGRpcShutdownDeadline(TDuration::Seconds(1)); + setup3.Start(); + setup1.Start(); + + // Check that consumer reads from all sources. + setup1.WriteToTopic({"msg1", "msg2"}); + setup2.WriteToTopic({"msg3", "msg4"}); + setup3.WriteToTopic({"msg10", "msg11", "msg12", "msg13", "msg14", "msg15"}); + auto consumer = setup1.StartConsumer(MakeConsumerSettings(setup1)); + setup1.ReadFromTopic({{"msg1", "msg2"}, {"msg3", "msg4"}, {"msg10", "msg11", "msg12", "msg13", "msg14", "msg15"}}, true, consumer.Get()); + + // Turn Dc off and check that consumer works. + setup1.GetLog() << TLOG_INFO << "Turn off DataCenter dc3"; + setup3.ShutdownGRpc(); + setup1.WriteToTopic({"a", "b"}); + setup2.WriteToTopic({"x", "y", "z"}); + setup1.ReadFromTopic({{"a", "b"}, {"x", "y", "z"}}, true, consumer.Get()); + + // Turn Dc on again and check work again. + setup1.GetLog() << TLOG_INFO << "Turn on DataCenter dc3"; + setup3.EnableGRpc(); + setup1.WriteToTopic({"1", "2"}); + setup3.WriteToTopic({"42"}); + setup1.ReadFromTopic({{"1", "2"}, {"42"}}, true, consumer.Get()); + } +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_producer.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_producer.cpp new file mode 100644 index 0000000000..6f63b0ab4a --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_producer.cpp @@ -0,0 +1,423 @@ +#include "multicluster_producer.h" +#include "persqueue_p.h" + +#include <util/string/builder.h> + +#include <algorithm> + +namespace NPersQueue { + +TWeightIndex::TWeightIndex(size_t count) + : Weights(count) + , Index(count) + , IsEnabled(count) + , EnabledCnt(0) +{ + Y_VERIFY(count > 0); +} + +void TWeightIndex::SetWeight(size_t i, unsigned weight) { + const unsigned oldWeight = Weights[i]; + const int diff = static_cast<int>(weight) - static_cast<int>(oldWeight); + if (diff != 0) { + Weights[i] = weight; + if (IsEnabled[i]) { + UpdateIndex(i, diff); + } + } +} + +void TWeightIndex::Enable(size_t i) { + if (!IsEnabled[i]) { + ++EnabledCnt; + IsEnabled[i] = true; + UpdateIndex(i, static_cast<int>(Weights[i])); + } +} + +void TWeightIndex::Disable(size_t i) { + if (IsEnabled[i]) { + --EnabledCnt; + IsEnabled[i] = false; + UpdateIndex(i, -static_cast<int>(Weights[i])); + } +} + +size_t TWeightIndex::Choose(unsigned randomNumber) const { + Y_VERIFY(randomNumber < WeightsSum()); + Y_VERIFY(!Index.empty()); + const auto choice = std::upper_bound(Index.begin(), Index.end(), randomNumber); + Y_VERIFY(choice != Index.end()); + return choice - Index.begin(); +} + +void TWeightIndex::UpdateIndex(size_t i, int weightDiff) { + for (size_t j = i; j < Index.size(); ++j) { + Index[j] += weightDiff; + } +} + +static TProducerCommitResponse MakeErrorCommitResponse(const TString& description, TProducerSeqNo seqNo, const TData& data) { + TWriteResponse res; + res.MutableError()->SetDescription(description); + res.MutableError()->SetCode(NErrorCode::ERROR); + return TProducerCommitResponse{seqNo, data, std::move(res)}; +} + +TMultiClusterProducer::TMultiClusterProducer(const TMultiClusterProducerSettings& settings, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger) + : IProducerImpl(std::move(destroyEventRef), std::move(pqLib)) + , Settings(settings) + , DeadPromise(NThreading::NewPromise<TError>()) + , StartPromise(NThreading::NewPromise<TProducerCreateResponse>()) + , State(EState::Created) + , Logger(std::move(logger)) + , WeightIndex(Settings.ServerWeights.size()) + , FallbackWeightIndex(Settings.ServerWeights.size()) + , Subproducers(Settings.ServerWeights.size()) + , FallbackSubproducers(Settings.ServerWeights.size()) +{ + if (Settings.ServerWeights.size() <= 1) { + ythrow yexception() << "MultiCluster producer is working with several servers. You'd better to create retrying producer in this case"; + } + if (Settings.MinimumWorkingDcsCount == 0) { + ythrow yexception() << "MinimumWorkingDcsCount can't be zero. We want that you have at least something working!"; + } + if (Settings.SourceIdPrefix.empty()) { + ythrow yexception() << "SourceIdPrefix must be nonempty"; + } + + bool reconnectWarningPrinted = false; + for (TMultiClusterProducerSettings::TServerWeight& w : Settings.ServerWeights) { + if (!w.Weight) { + ythrow yexception() << "Server weight must be greater than zero"; + } + if (!w.ProducerSettings.ReconnectOnFailure) { + w.ProducerSettings.ReconnectOnFailure = true; + if (!reconnectWarningPrinted) { + reconnectWarningPrinted = true; + WARN_LOG("ReconnectOnFailure setting is off. Turning it on for multicluster producer.", Settings.SourceIdPrefix, ""); + } + ythrow yexception() << "You must enable ReconnectOnFailure setting"; + } + if (w.ProducerSettings.MaxAttempts == std::numeric_limits<unsigned>::max()) { + ythrow yexception() << "MaxAttempts must be a finite number"; + } + } + + for (size_t i = 0; i < Settings.ServerWeights.size(); ++i) { + WeightIndex.SetWeight(i, Settings.ServerWeights[i].Weight); + FallbackWeightIndex.SetWeight(i, Settings.ServerWeights[i].Weight); + } +} + +TMultiClusterProducer::~TMultiClusterProducer() { + Destroy("Destructor called"); +} + +void TMultiClusterProducer::StartSubproducer(TInstant deadline, size_t i, bool fallback) { + std::weak_ptr<TMultiClusterProducer> self(shared_from_this()); + auto& sp = GetSubproducers(fallback)[i]; + sp.Producer = CreateSubproducer(Settings.ServerWeights[i].ProducerSettings, i, fallback); + sp.StartFuture = sp.Producer->Start(deadline); + sp.DeadFuture = sp.Producer->IsDead(); + PQLib->Subscribe(sp.StartFuture, + this, + [self, i, fallback](const auto&) { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->OnProducerStarted(i, fallback); + } + }); + PQLib->Subscribe(sp.DeadFuture, + this, + [self, i, fallback](const auto&) { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->OnProducerDead(i, fallback); + } + }); +} + +void TMultiClusterProducer::OnNeedProducerRestart(size_t i, bool fallback) { + if (State == EState::Dead) { + return; + } + const auto deadline = TInstant::Now() + Settings.ServerWeights[i].ProducerSettings.StartSessionTimeout; + StartSubproducer(deadline, i, fallback); +} + +void TMultiClusterProducer::ScheduleProducerRestart(size_t i, bool fallback) { + std::weak_ptr<TMultiClusterProducer> self(shared_from_this()); + PQLib->GetScheduler().Schedule(Settings.ServerWeights[i].ProducerSettings.ReconnectionDelay, + this, + [self, i, fallback] { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->OnNeedProducerRestart(i, fallback); + } + }); +} + +void TMultiClusterProducer::StartSubproducers(TInstant deadline, bool fallback) { + for (size_t i = 0; i < Settings.ServerWeights.size(); ++i) { + StartSubproducer(deadline, i, fallback); + } +} + +NThreading::TFuture<TProducerCreateResponse> TMultiClusterProducer::Start(TInstant deadline) noexcept { + Y_VERIFY(State == EState::Created); + State = EState::Starting; + StartSubproducers(deadline); + StartSubproducers(deadline, true); + return StartPromise.GetFuture(); +} + +void TMultiClusterProducer::OnProducerStarted(size_t i, bool fallback) { + if (State == EState::Dead) { + return; + } + TSubproducerInfo& info = GetSubproducers(fallback)[i]; + info.StartAnswered = true; + DEBUG_LOG("Start subproducer response: " << info.StartFuture.GetValue().Response, Settings.SourceIdPrefix, ""); + if (info.StartFuture.GetValue().Response.HasError()) { + ERR_LOG("Failed to start subproducer[" << i << "]: " << info.StartFuture.GetValue().Response.GetError().GetDescription(), Settings.SourceIdPrefix, ""); + if (State == EState::Starting && !fallback) { + size_t startedCount = 0; + for (const auto& p : Subproducers) { + startedCount += p.StartAnswered; + } + if (startedCount == Subproducers.size() && WeightIndex.EnabledCount() < Settings.MinimumWorkingDcsCount) { + Destroy("Not enough subproducers started successfully"); + return; + } + } + + ScheduleProducerRestart(i, fallback); + } else { + GetWeightIndex(fallback).Enable(i); + if (State == EState::Starting && !fallback && WeightIndex.EnabledCount() >= Settings.MinimumWorkingDcsCount) { + State = EState::Working; + TWriteResponse response; + response.MutableInit(); + StartPromise.SetValue(TProducerCreateResponse(std::move(response))); + } + if (fallback) { + ResendLastHopeQueue(i); + } + } +} + +void TMultiClusterProducer::OnProducerDead(size_t i, bool fallback) { + if (State == EState::Dead) { + return; + } + + WARN_LOG("Subproducer[" << i << "] is dead: " << GetSubproducers(fallback)[i].DeadFuture.GetValue().GetDescription(), Settings.SourceIdPrefix, ""); + GetWeightIndex(fallback).Disable(i); + if (!fallback && WeightIndex.EnabledCount() < Settings.MinimumWorkingDcsCount) { + Destroy("Not enough subproducers is online"); + return; + } + ScheduleProducerRestart(i, fallback); +} + +void TMultiClusterProducer::Destroy(const TString& description) { + TError error; + error.SetCode(NErrorCode::ERROR); + error.SetDescription(description); + Destroy(error); +} + +void TMultiClusterProducer::DestroyWrites(const TError& error, TIntrusiveListWithAutoDelete<TWriteInfo, TDelete>& pendingWrites) { + for (TWriteInfo& wi : pendingWrites) { + TWriteResponse response; + response.MutableError()->CopyFrom(error); + wi.ResponsePromise.SetValue(TProducerCommitResponse(0, std::move(wi.Data), std::move(response))); + } + pendingWrites.Clear(); +} + +void TMultiClusterProducer::Destroy(const TError& error) { + if (State == EState::Dead) { + return; + } + + DEBUG_LOG("Destroying multicluster producer. Reason: " << error, Settings.SourceIdPrefix, ""); + const EState prevState = State; + State = EState::Dead; + SubscribeDestroyed(); + if (prevState == EState::Starting) { + TWriteResponse response; + response.MutableError()->CopyFrom(error); + StartPromise.SetValue(TProducerCreateResponse(std::move(response))); + } + for (size_t i = 0; i < Subproducers.size(); ++i) { + Subproducers[i].Producer = nullptr; + FallbackSubproducers[i].Producer = nullptr; + DestroyWrites(error, FallbackSubproducers[i].LastHopePendingWrites); + } + DestroyWrites(error, PendingWrites); + DeadPromise.SetValue(error); + + DestroyPQLibRef(); +} + +void TMultiClusterProducer::Write(NThreading::TPromise<TProducerCommitResponse>& promise, TProducerSeqNo seqNo, TData data) noexcept { + return promise.SetValue(MakeErrorCommitResponse("MultiCluster producer doesn't allow to write with explicitly specified seq no.", seqNo, data)); +} + +void TMultiClusterProducer::Write(NThreading::TPromise<TProducerCommitResponse>& promise, TData data) noexcept { + Y_VERIFY(data.IsEncoded()); + + if (State != EState::Working) { + TWriteResponse res; + res.MutableError()->SetDescription(TStringBuilder() << "producer is " << (State == EState::Dead ? "dead" : "not ready") << ". Marker# PQLib01"); + res.MutableError()->SetCode(NErrorCode::ERROR); + promise.SetValue(TProducerCommitResponse{0, std::move(data), std::move(res)}); + return; + } + + TWriteInfo* wi = new TWriteInfo(); + PendingWrites.PushBack(wi); + + const size_t i = WeightIndex.RandomChoose(); + wi->Data = std::move(data); + wi->ResponsePromise = promise; + DelegateWrite(wi, i); +} + +void TMultiClusterProducer::DelegateWrite(TWriteInfo* wi, size_t i, bool fallback) { + if (wi->History.empty() || wi->History.back() != i) { + wi->History.push_back(i); + } + DEBUG_LOG("Delegate" << (fallback ? " fallback" : "") << " write to subproducer[" << i << "]", Settings.SourceIdPrefix, ""); + wi->ResponseFuture = GetSubproducers(fallback)[i].Producer->Write(wi->Data); + wi->WaitingCallback = true; + + std::weak_ptr<TMultiClusterProducer> self(shared_from_this()); + PQLib->Subscribe(wi->ResponseFuture, + this, + [self, wi](const auto&) { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->OnWriteResponse(wi); + } + }); +} + +void TMultiClusterProducer::ResendLastHopeQueue(size_t i) { + TSubproducerInfo& info = FallbackSubproducers[i]; + if (!info.LastHopePendingWrites.Empty()) { + for (TWriteInfo& wi : info.LastHopePendingWrites) { + if (!wi.WaitingCallback) { + DelegateWrite(&wi, i, true); + } + } + } +} + +void TMultiClusterProducer::OnWriteResponse(TWriteInfo* wi) { + if (State == EState::Dead) { + return; + } + + Y_VERIFY(wi->WaitingCallback); + Y_VERIFY(!wi->History.empty()); + + DEBUG_LOG("Write response: " << wi->ResponseFuture.GetValue().Response, Settings.SourceIdPrefix, ""); + wi->WaitingCallback = false; + if (wi->ResponseFuture.GetValue().Response.HasAck()) { + wi->ResponsePromise.SetValue(wi->ResponseFuture.GetValue()); + wi->Unlink(); + delete wi; + } else if (wi->ResponseFuture.GetValue().Response.HasError()) { + // retry write to other dc + if (wi->LastHope) { + const size_t i = wi->History.back(); + if (FallbackWeightIndex.Enabled(i)) { + DelegateWrite(wi, i, true); + } // else we will retry when this producer will become online. + return; + } + TWeightIndex index = FallbackWeightIndex; + for (size_t i : wi->History) { + index.Disable(i); + } + + // choose new fallback dc + size_t i; + if (index.EnabledCount() == 0) { + i = wi->History.back(); + } else { + i = index.RandomChoose(); + } + + // define whether this dc is the last hope + const bool lastHope = index.EnabledCount() <= 1; + if (lastHope) { + wi->LastHope = true; + wi->Unlink(); + FallbackSubproducers[i].LastHopePendingWrites.PushBack(wi); + } + + DEBUG_LOG("Retrying write to subproducer[" << i << "], lastHope: " << lastHope, Settings.SourceIdPrefix, ""); + DelegateWrite(wi, i, true); + } else { // incorrect structure of response + Y_VERIFY(false); + } +} + +NThreading::TFuture<TError> TMultiClusterProducer::IsDead() noexcept { + return DeadPromise.GetFuture(); +} + +TString TMultiClusterProducer::MakeSourceId(const TStringBuf& prefix, const TStringBuf& srcSourceId, + size_t producerIndex, bool fallback) { + TStringBuilder ret; + ret << prefix << "-" << producerIndex; + if (!srcSourceId.empty()) { + ret << "-"sv << srcSourceId; + } + if (fallback) { + ret << "-multicluster-fallback"sv; + } else { + ret << "-multicluster"sv; + } + return std::move(ret); +} + +std::shared_ptr<IProducerImpl> TMultiClusterProducer::CreateSubproducer(const TProducerSettings& srcSettings, + size_t index, bool fallback) { + TProducerSettings settings = srcSettings; + settings.SourceId = MakeSourceId(Settings.SourceIdPrefix, srcSettings.SourceId, index, fallback); + return PQLib->CreateRawRetryingProducer(settings, DestroyEventRef, Logger); +} + +void TMultiClusterProducer::SubscribeDestroyed() { + NThreading::TPromise<void> promise = ProducersDestroyed; + auto handler = [promise](const auto&) mutable { + promise.SetValue(); + }; + std::vector<NThreading::TFuture<void>> futures; + futures.reserve(Subproducers.size() * 2 + 1); + futures.push_back(DestroyedPromise.GetFuture()); + for (size_t i = 0; i < Subproducers.size(); ++i) { + if (Subproducers[i].Producer) { + futures.push_back(Subproducers[i].Producer->Destroyed()); + } + if (FallbackSubproducers[i].Producer) { + futures.push_back(FallbackSubproducers[i].Producer->Destroyed()); + } + } + WaitExceptionOrAll(futures).Subscribe(handler); +} + +NThreading::TFuture<void> TMultiClusterProducer::Destroyed() noexcept { + return ProducersDestroyed.GetFuture(); +} + +void TMultiClusterProducer::Cancel() { + Destroy(GetCancelReason()); +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_producer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_producer.h new file mode 100644 index 0000000000..16fcfb6b11 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_producer.h @@ -0,0 +1,147 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include "internals.h" +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iproducer.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> +#include "channel.h" +#include "iproducer_p.h" + +#include <library/cpp/threading/future/future.h> + +#include <util/generic/intrlist.h> +#include <util/random/random.h> + +#include <deque> + +namespace NPersQueue { + +// Structure to perform choose by weights +class TWeightIndex { +public: + // Creates zero-weighted disabled index + explicit TWeightIndex(size_t count); + + void SetWeight(size_t i, unsigned weight); + + void Enable(size_t i); + void Disable(size_t i); + + // Chooses by random number in interval [0, WeightsSum()) + size_t Choose(unsigned randomNumber) const; + + size_t RandomChoose() const { + return Choose(RandomNumber<unsigned>(WeightsSum())); + } + + size_t EnabledCount() const { + return EnabledCnt; + } + + bool Enabled(size_t i) const { + return i < IsEnabled.size() ? IsEnabled[i] : false; + } + + unsigned WeightsSum() const { + return Index.back(); + } + +private: + void UpdateIndex(size_t i, int weightDiff); + +private: + std::vector<unsigned> Weights; + std::vector<unsigned> Index; // structure for choosing. Index[i] is sum of enabled weights with index <= i. + std::vector<bool> IsEnabled; + size_t EnabledCnt; +}; + +class TMultiClusterProducer: public IProducerImpl, public std::enable_shared_from_this<TMultiClusterProducer> { +public: + TMultiClusterProducer(const TMultiClusterProducerSettings& settings, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger); + ~TMultiClusterProducer(); + + NThreading::TFuture<TProducerCreateResponse> Start(TInstant deadline) noexcept override; + + using IProducerImpl::Write; + + // Forbidden write + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TProducerSeqNo seqNo, TData data) noexcept override; + + // Allowed write + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TData data) noexcept override; + + NThreading::TFuture<TError> IsDead() noexcept override; + + NThreading::TFuture<void> Destroyed() noexcept override; + + void Cancel() override; + +private: + struct TWriteInfo: public TIntrusiveListItem<TWriteInfo> { + NThreading::TPromise<TProducerCommitResponse> ResponsePromise; + NThreading::TFuture<TProducerCommitResponse> ResponseFuture; + bool WaitingCallback = false; + std::vector<size_t> History; // Number of DCs with failed or pending writes + TData Data; + bool LastHope = false; // This DC is the last hope: retry till success. + // If this flag is set, this write info is in the local subproducer queue. + }; + + struct TSubproducerInfo { + size_t Number = std::numeric_limits<size_t>::max(); + std::shared_ptr<IProducerImpl> Producer; + NThreading::TFuture<TProducerCreateResponse> StartFuture; + NThreading::TFuture<TError> DeadFuture; + TIntrusiveListWithAutoDelete<TWriteInfo, TDelete> LastHopePendingWrites; // Local queue with writes only to this DC. + bool StartAnswered = false; + }; + +private: + std::shared_ptr<IProducerImpl> CreateSubproducer(const TProducerSettings& settings, size_t index, bool fallback = false); + + std::vector<TSubproducerInfo>& GetSubproducers(bool fallback = false) { + return fallback ? FallbackSubproducers : Subproducers; + } + + TWeightIndex& GetWeightIndex(bool fallback = false) { + return fallback ? FallbackWeightIndex : WeightIndex; + } + + static TString MakeSourceId(const TStringBuf& prefix, const TStringBuf& srcSourceId, + size_t producerIndex, bool fallback); + void StartSubproducers(TInstant deadline, bool fallback = false); + void StartSubproducer(TInstant deadline, size_t i, bool fallback = false); + void OnProducerStarted(size_t i, bool fallback); + void OnProducerDead(size_t i, bool fallback); + void OnWriteResponse(TWriteInfo* wi); + void ScheduleProducerRestart(size_t i, bool fallback); + void OnNeedProducerRestart(size_t i, bool fallback); // scheduler action + void Destroy(const TString& description); + void Destroy(const TError& error); + void DestroyWrites(const TError& error, TIntrusiveListWithAutoDelete<TWriteInfo, TDelete>& pendingWrites); + void DelegateWrite(TWriteInfo* wi, size_t i, bool fallback = false); + void ResendLastHopeQueue(size_t i); + void SubscribeDestroyed(); + +protected: + TMultiClusterProducerSettings Settings; + NThreading::TPromise<TError> DeadPromise; + NThreading::TPromise<TProducerCreateResponse> StartPromise; + enum class EState { + Created, + Starting, + Working, + Dead, + }; + EState State; + TIntrusivePtr<ILogger> Logger; + TWeightIndex WeightIndex; + TWeightIndex FallbackWeightIndex; + std::vector<TSubproducerInfo> Subproducers; + std::vector<TSubproducerInfo> FallbackSubproducers; + TIntrusiveListWithAutoDelete<TWriteInfo, TDelete> PendingWrites; + NThreading::TPromise<void> ProducersDestroyed = NThreading::NewPromise<void>(); +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_producer_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_producer_ut.cpp new file mode 100644 index 0000000000..6942681e6a --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/multicluster_producer_ut.cpp @@ -0,0 +1,360 @@ +#include "multicluster_producer.h" +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h> + +#include <library/cpp/testing/unittest/registar.h> + +#include <util/system/event.h> + +#include <atomic> + +namespace NPersQueue { + +Y_UNIT_TEST_SUITE(TWeightIndexTest) { + Y_UNIT_TEST(Enables) { + TWeightIndex wi(3); + UNIT_ASSERT_VALUES_EQUAL(wi.EnabledCount(), 0); + + wi.SetWeight(0, 1); + wi.SetWeight(1, 2); + UNIT_ASSERT_VALUES_EQUAL(wi.EnabledCount(), 0); + UNIT_ASSERT_VALUES_EQUAL(wi.WeightsSum(), 0); + + wi.Disable(1); + UNIT_ASSERT_VALUES_EQUAL(wi.EnabledCount(), 0); + UNIT_ASSERT_VALUES_EQUAL(wi.WeightsSum(), 0); + + wi.Enable(0); + wi.Enable(0); // no changes + UNIT_ASSERT_VALUES_EQUAL(wi.EnabledCount(), 1); + UNIT_ASSERT_VALUES_EQUAL(wi.WeightsSum(), 1); + + wi.Enable(1); + UNIT_ASSERT_VALUES_EQUAL(wi.EnabledCount(), 2); + UNIT_ASSERT_VALUES_EQUAL(wi.WeightsSum(), 3); + + wi.Disable(1); + UNIT_ASSERT_VALUES_EQUAL(wi.EnabledCount(), 1); + UNIT_ASSERT_VALUES_EQUAL(wi.WeightsSum(), 1); + + wi.Disable(1); // no changes + UNIT_ASSERT_VALUES_EQUAL(wi.EnabledCount(), 1); + UNIT_ASSERT_VALUES_EQUAL(wi.WeightsSum(), 1); + } + + Y_UNIT_TEST(Chooses) { + const size_t weightsCount = 3; + TWeightIndex wi(weightsCount); + UNIT_ASSERT_VALUES_EQUAL(wi.EnabledCount(), 0); + unsigned weights[weightsCount] = {1, 2, 3}; + for (size_t i = 0; i < weightsCount; ++i) { + wi.SetWeight(i, weights[i]); + wi.Enable(i); + } + UNIT_ASSERT_VALUES_EQUAL(wi.EnabledCount(), weightsCount); + UNIT_ASSERT_VALUES_EQUAL(wi.WeightsSum(), 6); + unsigned cnt[weightsCount] = {}; + for (size_t i = 0; i < wi.WeightsSum(); ++i) { + const size_t choice = wi.Choose(i); + UNIT_ASSERT(choice < weightsCount); // it is index + ++cnt[choice]; + } + for (size_t i = 0; i < weightsCount; ++i) { + UNIT_ASSERT_VALUES_EQUAL_C(weights[i], cnt[i], "i: " << i); + } + } +} + +Y_UNIT_TEST_SUITE(TMultiClusterProducerTest) { + Y_UNIT_TEST(NotStartedProducerCanBeDestructed) { + // Test that producer doesn't hang on till shutdown + TPQLib lib; + TMultiClusterProducerSettings settings; + settings.SourceIdPrefix = "prefix"; + settings.ServerWeights.resize(2); + + { + auto& w = settings.ServerWeights[0]; + w.Weight = 100500; + auto& s = w.ProducerSettings; + s.ReconnectOnFailure = true; + s.Server = TServerSetting{"localhost"}; + s.Topic = "topic1"; + s.MaxAttempts = 10; + } + + { + auto& w = settings.ServerWeights[1]; + w.Weight = 1; + auto& s = w.ProducerSettings; + s.ReconnectOnFailure = true; + s.Server = TServerSetting{"localhost"}; + s.Topic = "topic2"; + s.MaxAttempts = 10; + } + lib.CreateMultiClusterProducer(settings, {}, false); + } + + static void InitServer(TTestServer& testServer) { + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(false); + + testServer.AnnoyingClient->FullInit(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", 2); + + testServer.WaitInit("topic1"); + } + + static void ConnectsToMinOnlineDcs(size_t min) { + return; // Test is ignored. FIX: KIKIMR-7886 + + TTestServer testServer(false); + InitServer(testServer); + + TMultiClusterProducerSettings settings; + settings.SourceIdPrefix = "prefix"; + settings.ServerWeights.resize(2); + settings.MinimumWorkingDcsCount = min; + + { + auto& w = settings.ServerWeights[0]; + w.Weight = 100500; + auto& s = w.ProducerSettings; + s.ReconnectOnFailure = true; + s.Server = TServerSetting{"localhost", testServer.PortManager->GetPort()}; // port with nothing up + s.Topic = "topic1"; + s.ReconnectionDelay = TDuration::MilliSeconds(1); + s.MaxAttempts = 10; + } + + { + auto& w = settings.ServerWeights[1]; + w.Weight = 1; + auto& s = w.ProducerSettings; + s.ReconnectOnFailure = true; + s.Server = TServerSetting{"localhost", testServer.GrpcPort}; + s.Topic = "topic1"; + s.MaxAttempts = 10; + } + + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLib lib; + auto producer = lib.CreateMultiClusterProducer(settings, logger, false); + auto start = producer->Start(); + + if (min == 1) { + UNIT_ASSERT_C(!start.GetValueSync().Response.HasError(), start.GetValueSync().Response); + } else { + UNIT_ASSERT_C(start.GetValueSync().Response.HasError(), start.GetValueSync().Response); + } + + DestroyAndWait(producer); + } + + Y_UNIT_TEST(ConnectsToMinOnlineDcs) { + ConnectsToMinOnlineDcs(1); + } + + Y_UNIT_TEST(FailsToConnectToMinOnlineDcs) { + return; // Test is ignored. FIX: KIKIMR-7886 + + ConnectsToMinOnlineDcs(2); + } + + TConsumerSettings MakeConsumerSettings(const TTestServer& testServer) { + TConsumerSettings consumerSettings; + consumerSettings.ClientId = "user"; + consumerSettings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + consumerSettings.Topics.push_back("topic1"); + consumerSettings.CommitsDisabled = true; + return consumerSettings; + } + + void AssertWriteValid(const NThreading::TFuture<TProducerCommitResponse>& respFuture) { + const TProducerCommitResponse& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TWriteResponse::kAck, "Msg: " << resp.Response); + } + + void AssertReadValid(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kData, "Msg: " << resp.Response); + } + + std::pair<size_t, size_t> ReadNMessagesAndDestroyConsumers(THolder<IConsumer> consumer1, THolder<IConsumer> consumer2, size_t n) { + std::atomic<size_t> readsDone(0), reads1(0), reads2(0); + TSystemEvent ev; + std::pair<std::atomic<size_t>*, IConsumer*> infos[] = { + std::make_pair(&reads1, consumer1.Get()), + std::make_pair(&reads2, consumer2.Get()) + }; + for (size_t i = 0; i < n; ++i) { + for (auto& info : infos) { + auto* reads = info.first; + info.second->GetNextMessage().Subscribe([&, reads](const NThreading::TFuture<TConsumerMessage>& respFuture) { + ++(*reads); + const size_t newReadsDone = ++readsDone; + if (newReadsDone == n) { + ev.Signal(); + } + if (newReadsDone <= n) { + AssertReadValid(respFuture); + } + }); + } + } + ev.Wait(); + consumer1 = nullptr; + consumer2 = nullptr; + return std::pair<size_t, size_t>(reads1, reads2); + } + + Y_UNIT_TEST(WritesInSubproducers) { + return; // Test is ignored. FIX: KIKIMR-7886 + + TTestServer testServer1(false); + InitServer(testServer1); + + TTestServer testServer2(false); + InitServer(testServer2); + + TMultiClusterProducerSettings settings; + settings.SourceIdPrefix = "producer"; + settings.ServerWeights.resize(2); + settings.MinimumWorkingDcsCount = 1; + + { + auto& w = settings.ServerWeights[0]; + w.Weight = 1; + auto& s = w.ProducerSettings; + s.ReconnectOnFailure = true; + s.Server = TServerSetting{"localhost", testServer1.GrpcPort}; + s.Topic = "topic1"; + s.SourceId = "src0"; + s.MaxAttempts = 5; + s.ReconnectionDelay = TDuration::MilliSeconds(30); + } + + { + auto& w = settings.ServerWeights[1]; + w.Weight = 1; + auto& s = w.ProducerSettings; + s.ReconnectOnFailure = true; + s.Server = TServerSetting{"localhost", testServer2.GrpcPort}; + s.Topic = "topic1"; + s.SourceId = "src1"; + s.MaxAttempts = 2; + s.ReconnectionDelay = TDuration::MilliSeconds(30); + } + + TPQLibSettings pqLibSettings; + pqLibSettings.ThreadsCount = 1; + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLib lib(pqLibSettings); + auto producer = lib.CreateMultiClusterProducer(settings, logger, false); + auto startProducerResult = producer->Start(); + UNIT_ASSERT_C(!startProducerResult.GetValueSync().Response.HasError(), startProducerResult.GetValueSync().Response); + + auto consumer1 = lib.CreateConsumer(MakeConsumerSettings(testServer1), logger, false); + UNIT_ASSERT(!consumer1->Start().GetValueSync().Response.HasError()); + + auto consumer2 = lib.CreateConsumer(MakeConsumerSettings(testServer2), logger, false); + UNIT_ASSERT(!consumer2->Start().GetValueSync().Response.HasError()); + + const size_t writes = 100; + for (size_t i = 0; i < writes; ++i) { + auto write = producer->Write(TString(TStringBuilder() << "blob_1_" << i)); + AssertWriteValid(write); + } + + const std::pair<size_t, size_t> msgs = ReadNMessagesAndDestroyConsumers(std::move(consumer1), std::move(consumer2), writes); + UNIT_ASSERT_C(msgs.first > 30, "msgs.first = " << msgs.first); + UNIT_ASSERT_C(msgs.second > 30, "msgs.second = " << msgs.second); + + consumer1 = lib.CreateConsumer(MakeConsumerSettings(testServer1), logger, false); + UNIT_ASSERT(!consumer1->Start().GetValueSync().Response.HasError()); + + testServer2.ShutdownServer(); + //testServer2.CleverServer->ShutdownGRpc(); + + // Assert that all writes go in 1st DC + for (size_t i = 0; i < writes; ++i) { + auto write = producer->Write(TString(TStringBuilder() << "blob_2_" << i)); + auto read = consumer1->GetNextMessage(); + AssertWriteValid(write); + AssertReadValid(read); + } + + // Assert that all futures will be signalled after producer's death + std::vector<NThreading::TFuture<TProducerCommitResponse>> writesFutures(writes); + for (size_t i = 0; i < writes; ++i) { + writesFutures[i] = producer->Write(TString(TStringBuilder() << "blob_3_" << i)); + } + + auto isDead = producer->IsDead(); + producer = nullptr; + for (auto& f : writesFutures) { + f.GetValueSync(); + } + isDead.GetValueSync(); + + DestroyAndWait(consumer1); + DestroyAndWait(consumer2); + } + + Y_UNIT_TEST(CancelsOperationsAfterPQLibDeath) { + return; // Test is ignored. FIX: KIKIMR-7886 + + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(); + + const size_t partitions = 1; + testServer.AnnoyingClient->InitRoot(); + testServer.AnnoyingClient->InitDCs(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", partitions); + testServer.AnnoyingClient->InitSourceIds(); + + testServer.WaitInit("topic1"); + + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLibSettings pqLibSettings; + pqLibSettings.DefaultLogger = logger; + THolder<TPQLib> PQLib = MakeHolder<TPQLib>(pqLibSettings); + + TMultiClusterProducerSettings settings; + settings.SourceIdPrefix = "producer"; + settings.ServerWeights.resize(2); + settings.MinimumWorkingDcsCount = 1; + + for (auto& w : settings.ServerWeights) { + w.Weight = 1; + auto& s = w.ProducerSettings; + s.ReconnectOnFailure = true; + s.Server = TServerSetting{"localhost", testServer.GrpcPort}; + s.Topic = "topic1"; + s.SourceId = "src0"; + s.MaxAttempts = 5; + s.ReconnectionDelay = TDuration::MilliSeconds(30); + } + + auto producer = PQLib->CreateMultiClusterProducer(settings, logger, false); + UNIT_ASSERT(!producer->Start().GetValueSync().Response.HasError()); + auto isDead = producer->IsDead(); + UNIT_ASSERT(!isDead.HasValue()); + + auto write1 = producer->Write(1, TString("blob1")); + auto write2 = producer->Write(2, TString("blob2")); + + PQLib = nullptr; + + UNIT_ASSERT(write1.HasValue()); + UNIT_ASSERT(write2.HasValue()); + + auto write3 = producer->Write(3, TString("blob3")); + UNIT_ASSERT(write3.HasValue()); + UNIT_ASSERT(write3.GetValue().Response.HasError()); + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/persqueue.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/persqueue.cpp new file mode 100644 index 0000000000..cc1f46d835 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/persqueue.cpp @@ -0,0 +1,540 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include "internals.h" +#include "persqueue_p.h" +#include "compat_producer.h" +#include "consumer.h" +#include "multicluster_consumer.h" +#include "producer.h" +#include "retrying_consumer.h" +#include "retrying_producer.h" +#include "multicluster_producer.h" +#include "processor.h" +#include "compressing_producer.h" +#include "decompressing_consumer.h" +#include "local_caller.h" +#include "ydb_sdk_consumer.h" + +#include <util/digest/numeric.h> +#include <util/system/thread.h> + +namespace NPersQueue { + + +TString TPQLib::GetUserAgent() const { + return Impl->GetUserAgent(); +} + +void TPQLib::SetUserAgent(const TString& userAgent) { + Impl->SetUserAgent(userAgent); +} + +TPQLib::TPQLib(const TPQLibSettings& settings) + : Impl(new TPQLibPrivate(settings)) +{} + +TPQLib::~TPQLib() { + AtomicSet(Alive, 0); + Impl->CancelObjectsAndWait(); +} + +#define CHECK_PQLIB_ALIVE() Y_VERIFY(AtomicGet(Alive), "Attempt to use PQLib after/during its destruction.") + +THolder<IProducer> TPQLib::CreateMultiClusterProducer(const TMultiClusterProducerSettings& settings, TIntrusivePtr<ILogger> logger, bool deprecated) { + CHECK_PQLIB_ALIVE(); + return MakeHolder<TPublicProducer>(Impl->CreateMultiClusterProducer(settings, deprecated, std::move(logger))); +} + +THolder<IProducer> TPQLib::CreateProducer(const TProducerSettings& settings, TIntrusivePtr<ILogger> logger, bool deprecated) { + CHECK_PQLIB_ALIVE(); + auto producer = settings.ReconnectOnFailure ? + Impl->CreateRetryingProducer(settings, deprecated, std::move(logger)) : + Impl->CreateProducer(settings, deprecated, std::move(logger)); + return MakeHolder<TPublicProducer>(std::move(producer)); +} + +template <class TConsumerOrProducerSettings> +NYdb::NPersQueue::TPersQueueClient& TPQLibPrivate::GetOrCreatePersQueueClient(const TConsumerOrProducerSettings& settings, const TIntrusivePtr<ILogger>& logger) { + with_lock(Lock) { + if (!CompressExecutor) { + CompressExecutor = NYdb::NPersQueue::CreateThreadPoolExecutorAdapter(CompressionPool); + HandlersExecutor = NYdb::NPersQueue::CreateThreadPoolExecutor(1); + } + TWrapperKey key{settings.Server, settings.CredentialsProvider}; + auto it = Wrappers.find(key); + if (Wrappers.end() == it) { + CreateWrapper(key, logger); + it = Wrappers.find(key); + } + return *it->second.PersQueueClient; + } +} + +std::shared_ptr<IConsumerImpl> TPQLibPrivate::CreateNewConsumer(const TConsumerSettings& settings, TIntrusivePtr<ILogger> logger) +{ + if (Settings.EnableGRpcV1 && settings.Server.UseLogbrokerCDS != EClusterDiscoveryUsageMode::DontUse && !settings.CommitsDisabled && settings.Unpack) { + auto ret = std::make_shared<TLocalConsumerImplCaller<TYdbSdkCompatibilityConsumer>>(settings, + GetSelfRefsAreDeadEvent(false), + this, + logger, + GetOrCreatePersQueueClient(settings, logger)); + ret->Init(); + AddToDestroySet(ret); + return ret; + } + return nullptr; +} + +THolder<IConsumer> TPQLib::CreateConsumer(const TConsumerSettings& settings, TIntrusivePtr<ILogger> logger, bool deprecated) { + CHECK_PQLIB_ALIVE(); + return MakeHolder<TPublicConsumer>(Impl->CreateConsumer(settings, deprecated, std::move(logger))); +} + +THolder<IProcessor> TPQLib::CreateProcessor(const TProcessorSettings& settings, TIntrusivePtr<ILogger> logger, bool deprecated) { + CHECK_PQLIB_ALIVE(); + return MakeHolder<TPublicProcessor>(Impl->CreateProcessor(settings, deprecated, std::move(logger))); +} + +void TPQLib::SetLogger(TIntrusivePtr<ILogger> logger) { + CHECK_PQLIB_ALIVE(); + Impl->SetLogger(std::move(logger)); +} + +#undef CHECK_PQLIB_ALIVE + +static void DoExecute(std::shared_ptr<grpc::CompletionQueue> cq, std::shared_ptr<std::atomic<bool>> shuttingDown) { + TThread::SetCurrentThreadName("pqlib_grpc_thr"); + void* tag; + bool ok; + while (cq->Next(&tag, &ok)) { + IQueueEvent* const ev(static_cast<IQueueEvent*>(tag)); + if (shuttingDown->load() || !ev->Execute(ok)) { + ev->DestroyRequest(); + } + } +} + +TPQLibPrivate::TPQLibPrivate(const TPQLibSettings& settings) + : Settings(settings) + , CQ(std::make_shared<grpc::CompletionQueue>()) + , Scheduler(MakeHolder<TScheduler>(this)) + , Logger(settings.DefaultLogger) +{ + Y_VERIFY(Settings.ThreadsCount > 0); + Y_VERIFY(Settings.GRpcThreads > 0); + Y_VERIFY(Settings.CompressionPoolThreads > 0); + CompressionPool = std::make_shared<TThreadPool>(); + CompressionPool->Start(Settings.CompressionPoolThreads); + QueuePool.Start(Settings.ThreadsCount); + GRpcThreads.reserve(Settings.GRpcThreads); + for (size_t i = 0; i < Settings.GRpcThreads; ++i) { + GRpcThreads.emplace_back([cq = CQ, shuttingDown = ShuttingDown](){ DoExecute(cq, shuttingDown); }); + } + CreateSelfRefsAreDeadPtr(); +} + +TPQLibPrivate::~TPQLibPrivate() { + DEBUG_LOG("Destroying PQLib. Destroying scheduler.", "", ""); + Scheduler.Reset(); + + DEBUG_LOG("Destroying PQLib. Set stop flag.", "", ""); + Stop(); + + if (CQ) { + DEBUG_LOG("Destroying PQLib. Shutting down completion queue.", "", ""); + CQ->Shutdown(); + } + + if (!AtomicGet(HasDeprecatedObjects)) { + DEBUG_LOG("Destroying PQLib. Joining threads.", "", ""); + for (std::thread& thread : GRpcThreads) { + thread.join(); + } + } else { + DEBUG_LOG("Destroying PQLib. Detaching threads.", "", ""); + for (std::thread& thread : GRpcThreads) { + thread.detach(); + } + } + + if (YdbDriver) { + YdbDriver->Stop(); + YdbDriver = nullptr; + } + + DEBUG_LOG("Destroying PQLib. Stopping compression pool.", "", ""); + CompressionPool->Stop(); + + DEBUG_LOG("Destroying PQLib. Stopping queue pool.", "", ""); + QueuePool.Stop(); +} + +std::shared_ptr<IProducerImpl> TPQLibPrivate::CreateRawProducer(const TProducerSettings& settings, const TChannelInfo& channelInfo, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger) +{ + std::shared_ptr<IProducerImpl> res; + auto rawProducerImpl = CreateRawProducerImpl<TProducer>(settings, refsAreDeadPtr, ChooseLogger(logger)); + rawProducerImpl->SetChannel(channelInfo); + res = std::move(rawProducerImpl); + return res; +} + +std::shared_ptr<IProducerImpl> TPQLibPrivate::CreateRawProducer(const TProducerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger) +{ + TChannelHolder t = CreateChannel(settings, ChooseLogger(logger)); + std::shared_ptr<IProducerImpl> res; + + auto rawProducerImpl = CreateRawProducerImpl<TProducer>(settings, refsAreDeadPtr, ChooseLogger(logger)); + rawProducerImpl->SetChannel(t); + res = std::move(rawProducerImpl); + return res; +} + +std::shared_ptr<IProducerImpl> TPQLibPrivate::CreateProducer(const TProducerSettings& settings, const TChannelInfo& channelInfo, bool deprecated, TIntrusivePtr<ILogger> logger) +{ + DEBUG_LOG("Create producer", settings.SourceId, ""); + std::shared_ptr<IProducerImpl> rawProducer = CreateRawProducer(settings, channelInfo, GetSelfRefsAreDeadEvent(deprecated), ChooseLogger(logger)); + std::shared_ptr<IProducerImpl> compressing = std::make_shared<TLocalProducerImplCaller<TCompressingProducer>>(std::move(rawProducer), settings.Codec, settings.Quality, GetSelfRefsAreDeadEvent(deprecated), this, ChooseLogger(logger)); + compressing->Init(); + if (!deprecated) { + AddToDestroySet(compressing); + } + return compressing; +} + +class TLoggerWrapper : public TLogBackend { +public: + TLoggerWrapper(TIntrusivePtr<ILogger> logger) noexcept + : Logger(logger) + {} + + ~TLoggerWrapper() {} + + void WriteData(const TLogRecord& rec) + { + Logger->Log(TString(rec.Data, rec.Len), "", "", (int)rec.Priority); + } + + void ReopenLog() {} + +private: + TIntrusivePtr<ILogger> Logger; +}; + + +void TPQLibPrivate::CreateWrapper(const TWrapperKey& key, TIntrusivePtr<ILogger> logger) { + auto& wrapper = Wrappers[key]; + const auto& settings = key.Settings; + if (!YdbDriver) { + NYdb::TDriverConfig config; + + config.SetEndpoint(TStringBuilder() << settings.Address << ":" << settings.Port) //TODO BAD + .SetDatabase(settings.Database.empty() ? "/Root" : settings.Database) + .SetNetworkThreadsNum(Settings.GRpcThreads) + .SetClientThreadsNum(Settings.ThreadsCount) + .SetGRpcKeepAliveTimeout(TDuration::Seconds(90)) + .SetGRpcKeepAlivePermitWithoutCalls(true) + .SetDiscoveryMode(NYdb::EDiscoveryMode::Async); + if (logger) + config.SetLog(MakeHolder<TLoggerWrapper>(logger)); + if (settings.UseSecureConnection) + config.UseSecureConnection(settings.CaCert); + + YdbDriver = MakeHolder<NYdb::TDriver>(config); + } + + NYdb::NPersQueue::TPersQueueClientSettings pqSettings; + pqSettings.DefaultCompressionExecutor(CompressExecutor) + .DefaultHandlersExecutor(HandlersExecutor) + .ClusterDiscoveryMode(NYdb::NPersQueue::EClusterDiscoveryMode::Auto) + .DiscoveryEndpoint(TStringBuilder() << settings.Address << ":" << settings.Port) + .Database(settings.Database); + if (key.Provider) { + pqSettings.CredentialsProviderFactory(std::shared_ptr<NYdb::ICredentialsProviderFactory>(new TCredentialsProviderFactoryWrapper(key.Provider))); + } + wrapper.PersQueueClient = MakeHolder<NYdb::NPersQueue::TPersQueueClient>(*YdbDriver, pqSettings); +} + +std::shared_ptr<IProducerImpl> TPQLibPrivate::CreateNewProducer(const TProducerSettings& settings, TIntrusivePtr<ILogger> logger) +{ + if (Settings.EnableGRpcV1) { + auto ret = std::make_shared<TLocalProducerImplCaller<TYdbSdkCompatibilityProducer>>(settings, + GetOrCreatePersQueueClient(settings, logger), + GetSelfRefsAreDeadEvent(false), + this); + ret->Init(); + AddToDestroySet(ret); + return ret; + } + return nullptr; +} + +std::shared_ptr<IProducerImpl> TPQLibPrivate::CreateProducer(const TProducerSettings& settings, bool deprecated, TIntrusivePtr<ILogger> logger) +{ + auto newP = CreateNewProducer(settings, ChooseLogger(logger)); + if (newP) return newP; + + DEBUG_LOG("Create producer", settings.SourceId, ""); + std::shared_ptr<IProducerImpl> rawProducer = CreateRawProducer(settings, GetSelfRefsAreDeadEvent(deprecated), ChooseLogger(logger)); + std::shared_ptr<IProducerImpl> compressing = std::make_shared<TLocalProducerImplCaller<TCompressingProducer>>(std::move(rawProducer), settings.Codec, settings.Quality, GetSelfRefsAreDeadEvent(deprecated), this, ChooseLogger(logger)); + compressing->Init(); + if (!deprecated) { + AddToDestroySet(compressing); + } + return compressing; +} + +template<typename TRawProducerImpl> +std::shared_ptr<TRawProducerImpl> TPQLibPrivate::CreateRawProducerImpl(const TProducerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger) +{ + DEBUG_LOG("Create raw producer", settings.SourceId, ""); + const bool addToDestroySet = refsAreDeadPtr != nullptr; + NThreading::TPromise<TProducerCreateResponse> promise = NThreading::NewPromise<TProducerCreateResponse>(); + auto ret = std::make_shared<TLocalProducerImplCaller<TRawProducerImpl>>(settings, CQ, promise, std::move(refsAreDeadPtr), this, ChooseLogger(logger)); + ret->Init(); + if (addToDestroySet) { + AddToDestroySet(ret); + } + return ret; +} + +std::shared_ptr<IConsumerImpl> TPQLibPrivate::CreateRawConsumer(const TConsumerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger) +{ + auto t = CreateChannel(settings.Server, settings.CredentialsProvider, ChooseLogger(logger), settings.PreferLocalProxy); + auto res = CreateRawConsumerImpl(settings, std::move(refsAreDeadPtr), ChooseLogger(logger)); + res->SetChannel(t); + return std::move(res); +} + +std::shared_ptr<IConsumerImpl> TPQLibPrivate::CreateRawConsumer(const TConsumerSettings& settings, const TChannelInfo& channelInfo, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger) +{ + auto res = CreateRawConsumerImpl(settings, refsAreDeadPtr, ChooseLogger(logger)); + res->SetChannel(channelInfo); + return std::move(res); +} + +std::shared_ptr<IConsumerImpl> TPQLibPrivate::CreateRawMultiClusterConsumer(const TConsumerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger) +{ + DEBUG_LOG("Create multicluster consumer", "", ""); + const bool addToDestroySet = refsAreDeadPtr != nullptr; + auto consumer = std::make_shared<TLocalConsumerImplCaller<TMultiClusterConsumer>>(settings, std::move(refsAreDeadPtr), this, logger); + consumer->Init(); + if (addToDestroySet) { + AddToDestroySet(consumer); + } + return std::move(consumer); +} + +std::shared_ptr<IConsumerImpl> TPQLibPrivate::CreateDecompressingConsumer(std::shared_ptr<IConsumerImpl> subconsumer, const TConsumerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger) { + const bool addToDestroySet = refsAreDeadPtr != nullptr; + std::shared_ptr<IConsumerImpl> consumer = std::make_shared<TLocalConsumerImplCaller<TDecompressingConsumer>>(std::move(subconsumer), settings, std::move(refsAreDeadPtr), this, ChooseLogger(logger)); + consumer->Init(); + if (addToDestroySet) { + AddToDestroySet(consumer); + } + return consumer; +} + +std::shared_ptr<IConsumerImpl> TPQLibPrivate::CreateRawRetryingConsumer(const TConsumerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger) { + const bool addToDestroySet = refsAreDeadPtr != nullptr; + std::shared_ptr<IConsumerImpl> consumer = std::make_shared<TLocalConsumerImplCaller<TRetryingConsumer>>(settings, std::move(refsAreDeadPtr), this, logger); + consumer->Init(); + if (addToDestroySet) { + AddToDestroySet(consumer); + } + return consumer; +} + +std::shared_ptr<IConsumerImpl> TPQLibPrivate::CreateConsumer(const TConsumerSettings& settings, bool deprecated, TIntrusivePtr<ILogger> logger) +{ + if (auto newConsumer = CreateNewConsumer(settings, ChooseLogger(logger))) { + return newConsumer; + } + + DEBUG_LOG("Create consumer", "", ""); + // Create raw consumer. + std::shared_ptr<IConsumerImpl> consumer; + if (settings.ReadFromAllClusterSources) { + if (settings.ReadMirroredPartitions) { + ythrow yexception() << "Can't create consumer with both ReadMirroredPartitions and ReadFromAllClusterSources options"; + } + if (settings.Server.UseLogbrokerCDS == EClusterDiscoveryUsageMode::DontUse) { + ythrow yexception() << "Can't create consumer with ReadFromAllClusterSources option, but without using cluster discovery"; + } + consumer = CreateRawMultiClusterConsumer(settings, GetSelfRefsAreDeadEvent(deprecated), ChooseLogger(logger)); + } else if (settings.ReconnectOnFailure) { + consumer = CreateRawRetryingConsumer(settings, GetSelfRefsAreDeadEvent(deprecated), ChooseLogger(logger)); + } else { + consumer = CreateRawConsumer(settings, GetSelfRefsAreDeadEvent(deprecated), ChooseLogger(logger)); + } + if (settings.Unpack) { + consumer = CreateDecompressingConsumer(std::move(consumer), settings, GetSelfRefsAreDeadEvent(deprecated), ChooseLogger(logger)); + } + return consumer; +} + +std::shared_ptr<IConsumerImpl> TPQLibPrivate::CreateConsumer(const TConsumerSettings& settings, const TChannelInfo& channelInfo, bool deprecated, TIntrusivePtr<ILogger> logger) +{ + DEBUG_LOG("Create consumer", "", ""); + std::shared_ptr<IConsumerImpl> consumer = CreateRawConsumer(settings, channelInfo, GetSelfRefsAreDeadEvent(deprecated), ChooseLogger(logger)); + if (settings.Unpack) { + return CreateDecompressingConsumer(std::move(consumer), settings, GetSelfRefsAreDeadEvent(deprecated), ChooseLogger(logger)); + } + return consumer; +} + +std::shared_ptr<TConsumer> TPQLibPrivate::CreateRawConsumerImpl(const TConsumerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger) +{ + DEBUG_LOG("Create raw consumer", "", ""); + const bool addToDestroySet = refsAreDeadPtr != nullptr; + NThreading::TPromise<TConsumerCreateResponse> promise = NThreading::NewPromise<TConsumerCreateResponse>(); + std::shared_ptr<TConsumer> consumer = std::make_shared<TLocalConsumerImplCaller<TConsumer>>(settings, CQ, promise, std::move(refsAreDeadPtr), this, ChooseLogger(logger)); + consumer->Init(); + if (addToDestroySet) { + AddToDestroySet(consumer); + } + return consumer; +} + +std::shared_ptr<IProducerImpl> TPQLibPrivate::CreateRawRetryingProducer(const TProducerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger) +{ + DEBUG_LOG("Create raw retrying producer", settings.SourceId, ""); + const bool addToDestroySet = refsAreDeadPtr != nullptr; + auto producer = std::make_shared<TLocalProducerImplCaller<TRetryingProducer>>(settings, std::move(refsAreDeadPtr), this, ChooseLogger(logger)); + producer->Init(); + if (addToDestroySet) { + AddToDestroySet(producer); + } + return std::move(producer); +} + +std::shared_ptr<IProducerImpl> TPQLibPrivate::CreateRetryingProducer(const TProducerSettings& settings, bool deprecated, TIntrusivePtr<ILogger> logger) +{ + auto newP = CreateNewProducer(settings, ChooseLogger(logger)); + if (newP) return newP; + + DEBUG_LOG("Create retrying producer", settings.SourceId, ""); + std::shared_ptr<IProducerImpl> rawProducer = CreateRawRetryingProducer(settings, GetSelfRefsAreDeadEvent(deprecated), ChooseLogger(logger)); + std::shared_ptr<IProducerImpl> compressing = std::make_shared<TLocalProducerImplCaller<TCompressingProducer>>(std::move(rawProducer), settings.Codec, settings.Quality, GetSelfRefsAreDeadEvent(deprecated), this, ChooseLogger(logger)); + compressing->Init(); + if (!deprecated) { + AddToDestroySet(compressing); + } + return compressing; +} + +std::shared_ptr<IProducerImpl> TPQLibPrivate::CreateRawMultiClusterProducer(const TMultiClusterProducerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger) +{ + DEBUG_LOG("Create raw multicluster producer", settings.SourceIdPrefix, ""); + const bool addToDestroySet = refsAreDeadPtr != nullptr; + auto producer = std::make_shared<TLocalProducerImplCaller<TMultiClusterProducer>>(settings, std::move(refsAreDeadPtr), this, ChooseLogger(logger)); + producer->Init(); + if (addToDestroySet) { + AddToDestroySet(producer); + } + return std::move(producer); +} + +std::shared_ptr<IProducerImpl> TPQLibPrivate::CreateMultiClusterProducer(const TMultiClusterProducerSettings& settings, bool deprecated, TIntrusivePtr<ILogger> logger) +{ + DEBUG_LOG("Create multicluster producer", settings.SourceIdPrefix, ""); + if (settings.ServerWeights.empty()) { + ythrow yexception() << "Can't create producer without server"; + } + const ECodec codec = settings.ServerWeights[0].ProducerSettings.Codec; + const int quality = settings.ServerWeights[0].ProducerSettings.Quality; + for (const auto& w : settings.ServerWeights) { + if (w.ProducerSettings.Codec != codec) { + ythrow yexception() << "Can't create multicluster producer with different codecs"; + } + } + + std::shared_ptr<IProducerImpl> rawProducer = CreateRawMultiClusterProducer(settings, GetSelfRefsAreDeadEvent(deprecated), ChooseLogger(logger)); + std::shared_ptr<IProducerImpl> compressing = std::make_shared<TLocalProducerImplCaller<TCompressingProducer>>(std::move(rawProducer), codec, quality, GetSelfRefsAreDeadEvent(deprecated), this, ChooseLogger(logger)); + compressing->Init(); + if (!deprecated) { + AddToDestroySet(compressing); + } + return compressing; +} + +std::shared_ptr<IProcessorImpl> TPQLibPrivate::CreateProcessor(const TProcessorSettings& settings, bool deprecated, TIntrusivePtr<ILogger> logger) +{ + DEBUG_LOG("Create processor", "", ""); + auto processor = std::make_shared<TLocalProcessorImplCaller<TProcessor>>(settings, GetSelfRefsAreDeadEvent(deprecated), this, ChooseLogger(logger)); + processor->Init(); + if (!deprecated) { + AddToDestroySet(processor); + } + return processor; +} + +TChannelHolder TPQLibPrivate::CreateChannel( + const TServerSetting& server, const TCredProviderPtr& credentialsProvider, TIntrusivePtr<ILogger> logger, + bool preferLocalProxy +) { + DEBUG_LOG("Create channel", "", ""); + TChannelHolder res; + res.ChannelPtr = new TChannel(server, credentialsProvider, this, ChooseLogger(logger), preferLocalProxy); + res.ChannelInfo = res.ChannelPtr->GetChannel(); + res.ChannelPtr->Start(); + return res; +} + + +TChannelHolder TPQLibPrivate::CreateChannel( + const TProducerSettings& settings, TIntrusivePtr<ILogger> logger, bool preferLocalProxy +) { + DEBUG_LOG("Create channel", "", ""); + TChannelHolder res; + res.ChannelPtr = new TChannel(settings, this, ChooseLogger(logger), preferLocalProxy); + res.ChannelInfo = res.ChannelPtr->GetChannel(); + res.ChannelPtr->Start(); + return res; +} + +void TPQLibPrivate::CancelObjectsAndWait() { + { + auto guard = Guard(Lock); + DEBUG_LOG("Destroying PQLib. Cancelling objects: " << SyncDestroyedSet.size(), "", ""); + for (auto& ptr : SyncDestroyedSet) { + auto ref = ptr.lock(); + if (ref) { + ref->Cancel(); + } + } + SelfRefsAreDeadPtr = nullptr; + } + DEBUG_LOG("Destroying PQLib. Waiting refs to die", "", ""); + SelfRefsAreDeadEvent.Wait(); + + DEBUG_LOG("Destroying PQLib. HasDeprecatedObjects: " << AtomicGet(HasDeprecatedObjects) << ". Refs to PQLibPrivate: " << RefCount(), "", ""); + if (!AtomicGet(HasDeprecatedObjects)) { + Y_ASSERT(RefCount() == 1); + } +} + +void TPQLibPrivate::CreateSelfRefsAreDeadPtr() { + void* fakeAddress = &SelfRefsAreDeadEvent; + SelfRefsAreDeadPtr = std::shared_ptr<void>(fakeAddress, + [event = SelfRefsAreDeadEvent](void*) mutable { + event.Signal(); + }); +} + +void TPQLibPrivate::AddToDestroySet(std::weak_ptr<TSyncDestroyed> obj) { + auto guard = Guard(Lock); + // Check dead objects first + if (DestroySetAddCounter * 2 > SyncDestroyedSet.size()) { + for (auto i = SyncDestroyedSet.begin(); i != SyncDestroyedSet.end();) { + if (i->expired()) { + i = SyncDestroyedSet.erase(i); + } else { + ++i; + } + } + DestroySetAddCounter = 0; + } + ++DestroySetAddCounter; + SyncDestroyedSet.emplace(std::move(obj)); +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/persqueue_p.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/persqueue_p.h new file mode 100644 index 0000000000..4c7145b048 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/persqueue_p.h @@ -0,0 +1,316 @@ +#pragma once + +#include "channel.h" +#include "internals.h" +#include "iconsumer_p.h" +#include "iprocessor_p.h" +#include "iproducer_p.h" +#include "scheduler.h" +#include "queue_pool.h" + +#include <kikimr/public/sdk/cpp/client/ydb_persqueue/persqueue.h> + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iprocessor.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iproducer.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/iconsumer.h> + +#include <library/cpp/threading/future/future.h> + +#include <util/generic/ptr.h> +#include <util/generic/string.h> +#include <util/generic/vector.h> +#include <util/generic/maybe.h> +#include <util/generic/queue.h> +#include <util/generic/guid.h> +#include <util/generic/hash_set.h> +#include <util/system/types.h> +#include <util/system/mutex.h> +#include <util/thread/factory.h> +#include <util/thread/pool.h> + +#include <util/datetime/base.h> + +#include <grpc++/create_channel.h> +#include <grpc++/completion_queue.h> +#include <google/protobuf/message.h> + +#include <atomic> +#include <type_traits> +#include <memory> +#include <set> +#include <thread> + +namespace NPersQueue { + +class TProducer; +class TConsumer; + + +struct TWrapperKey { + TServerSetting Settings; + std::shared_ptr<ICredentialsProvider> Provider; + + bool operator <(const TWrapperKey& b) const { + return Settings < b.Settings || (!(b.Settings < Settings) && Provider < b.Provider); + } + +}; + + + +class TCredentialsProviderWrapper : public NYdb::ICredentialsProvider { +public: + TCredentialsProviderWrapper(std::shared_ptr<NPersQueue::ICredentialsProvider> provider) + : Provider(provider) + {} + + bool IsValid() const override { + return true; + } + + TString GetAuthInfo() const override { + return GetToken(Provider.get()); + } + + private: + std::shared_ptr<NPersQueue::ICredentialsProvider> Provider; +}; + + +class TCredentialsProviderFactoryWrapper : public NYdb::ICredentialsProviderFactory { +public: + TString GetClientIdentity() const override { + return Guid; + } + + std::shared_ptr<NYdb::ICredentialsProvider> CreateProvider() const override { + return Provider; + }; + + + TCredentialsProviderFactoryWrapper(std::shared_ptr<NPersQueue::ICredentialsProvider> provider) + : Provider(new TCredentialsProviderWrapper(provider)) + , Guid(CreateGuidAsString()) + {} + +private: + std::shared_ptr<NYdb::ICredentialsProvider> Provider; + TString Guid; +}; + +class TPQLibPrivate: public TAtomicRefCount<TPQLibPrivate> { + template <class T> + class TBaseFutureSubscriberObjectInQueue: public IObjectInQueue { + public: + void SetFuture(const NThreading::TFuture<T>& future) { + Future = future; + } + + protected: + NThreading::TFuture<T> Future; + }; + + template <class T, class TFunc> + class TFutureSubscriberObjectInQueue: public TBaseFutureSubscriberObjectInQueue<T> { + public: + template <class TF> + TFutureSubscriberObjectInQueue(TF&& f) + : Func(std::forward<TF>(f)) + { + } + + void Process(void*) override { + THolder<IObjectInQueue> holder(this); // like in TOwnedObjectInQueue::Process() + Func(this->Future); + } + + private: + TFunc Func; + }; + + template <class T, class TFunc> + static THolder<TBaseFutureSubscriberObjectInQueue<T>> MakeFutureSubscribeObjectInQueue(TFunc&& f) { + return MakeHolder<TFutureSubscriberObjectInQueue<T, std::remove_reference_t<TFunc>>>(std::forward<TFunc>(f)); + } + + template <class T> + class TSubscriberFunc { + public: + TSubscriberFunc(TPQLibPrivate* pqLib, const void* queueTag, THolder<TBaseFutureSubscriberObjectInQueue<T>> callback) + : PQLib(pqLib) + , QueueTag(queueTag) + , Callback(std::make_shared<THolder<TBaseFutureSubscriberObjectInQueue<T>>>(std::move(callback))) + { + } + + void operator()(const NThreading::TFuture<T>& future) { + if (Callback) { + (*Callback)->SetFuture(future); + THolder<IObjectInQueue> callback(std::move(*Callback)); + const bool added = PQLib->GetQueuePool().GetQueue(QueueTag).Add(callback.Get()); + if (added) { // can be false on shutdown + Y_UNUSED(callback.Release()); + } + } + } + + private: + TPQLibPrivate* PQLib; + const void* QueueTag; + // std::shared_ptr is only because copy constructor is required + std::shared_ptr<THolder<TBaseFutureSubscriberObjectInQueue<T>>> Callback; + }; + +public: + TPQLibPrivate(const TPQLibSettings& settings); + TPQLibPrivate(const TPQLibPrivate&) = delete; + TPQLibPrivate(TPQLibPrivate&&) = delete; + ~TPQLibPrivate(); + + TChannelHolder CreateChannel(const TServerSetting& server, const TCredProviderPtr& credentialsProvider, + TIntrusivePtr<ILogger> logger, bool preferLocalProxy = false); + TChannelHolder CreateChannel(const TProducerSettings& settings, TIntrusivePtr<ILogger> logger, + bool preferLocalProxy = false); + + + std::shared_ptr<IProducerImpl> CreateRawProducer(const TProducerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger); + std::shared_ptr<IProducerImpl> CreateRawProducer(const TProducerSettings& settings, const TChannelInfo& channelInfo, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger); + std::shared_ptr<IProducerImpl> CreateProducer(const TProducerSettings& settings, bool deprecated, TIntrusivePtr<ILogger> logger); + std::shared_ptr<IProducerImpl> CreateProducer(const TProducerSettings& settings, const TChannelInfo& channelInfo, bool deprecated, TIntrusivePtr<ILogger> logger); + + std::shared_ptr<IProducerImpl> CreateNewProducer(const TProducerSettings& settings, TIntrusivePtr<ILogger> logger); + + std::shared_ptr<IProducerImpl> CreateRawRetryingProducer(const TProducerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger); + std::shared_ptr<IProducerImpl> CreateRetryingProducer(const TProducerSettings& settings, bool deprecated, TIntrusivePtr<ILogger> logger); + + std::shared_ptr<IProducerImpl> CreateRawMultiClusterProducer(const TMultiClusterProducerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger); + std::shared_ptr<IProducerImpl> CreateMultiClusterProducer(const TMultiClusterProducerSettings& settings, bool deprecated, TIntrusivePtr<ILogger> logger); + + std::shared_ptr<IConsumerImpl> CreateRawConsumer(const TConsumerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger); + std::shared_ptr<IConsumerImpl> CreateRawConsumer(const TConsumerSettings& settings, const TChannelInfo& channelInfo, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger); + std::shared_ptr<IConsumerImpl> CreateRawRetryingConsumer(const TConsumerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger); + std::shared_ptr<IConsumerImpl> CreateRawMultiClusterConsumer(const TConsumerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger); + + std::shared_ptr<IConsumerImpl> CreateConsumer(const TConsumerSettings& settings, bool deprecated, TIntrusivePtr<ILogger> logger); + std::shared_ptr<IConsumerImpl> CreateConsumer(const TConsumerSettings& settings, const TChannelInfo& channelInfo, bool deprecated, TIntrusivePtr<ILogger> logger); + + std::shared_ptr<IConsumerImpl> CreateNewConsumer(const TConsumerSettings& settings, TIntrusivePtr<ILogger> logger); + + std::shared_ptr<IProcessorImpl> CreateProcessor(const TProcessorSettings& settings, bool deprecated, TIntrusivePtr<ILogger> logger); + + TString GetUserAgent() const { + auto guard = Guard(UALock); + return UserAgent; + } + + void SetUserAgent(const TString& userAgent) { + auto guard = Guard(UALock); + UserAgent = userAgent; + } + + void Stop() { + ShuttingDown->store(true); + } + + const std::shared_ptr<grpc::CompletionQueue>& GetCompletionQueue() const { + return CQ; + } + + const TPQLibSettings& GetSettings() const { + return Settings; + } + + TScheduler& GetScheduler() { + return *Scheduler; + } + + TQueuePool& GetQueuePool() { + return QueuePool; + } + + IThreadPool& GetCompressionPool() { + return *CompressionPool; + } + + template <class TFunc, class T> + void Subscribe(const NThreading::TFuture<T>& future, const void* queueTag, TFunc&& func) { + TSubscriberFunc<T> f(this, queueTag, MakeFutureSubscribeObjectInQueue<T>(std::forward<TFunc>(func))); + future.Subscribe(std::move(f)); + } + + void SetLogger(TIntrusivePtr<ILogger> logger) { + Logger = std::move(logger); + } + + TIntrusivePtr<ILogger> ChooseLogger(TIntrusivePtr<ILogger> logger) { + return logger ? std::move(logger) : Logger; + } + + void CancelObjectsAndWait(); + + // for testing + std::shared_ptr<void> GetSelfRefsAreDeadPtr() const { + return SelfRefsAreDeadPtr; + } + + void AddToDestroySet(std::weak_ptr<TSyncDestroyed> obj); + +private: + template<typename TRawProducerImpl> + std::shared_ptr<TRawProducerImpl> CreateRawProducerImpl(const TProducerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger); + std::shared_ptr<TConsumer> CreateRawConsumerImpl(const TConsumerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger); + std::shared_ptr<IConsumerImpl> CreateDecompressingConsumer(std::shared_ptr<IConsumerImpl> subconsumer, const TConsumerSettings& settings, std::shared_ptr<void> refsAreDeadPtr, TIntrusivePtr<ILogger> logger); + + void CreateWrapper(const TWrapperKey& key, TIntrusivePtr<ILogger> logger); + + template <class TConsumerOrProducerSettings> + NYdb::NPersQueue::TPersQueueClient& GetOrCreatePersQueueClient(const TConsumerOrProducerSettings& settings, const TIntrusivePtr<ILogger>& logger); + + std::shared_ptr<void> GetSelfRefsAreDeadEvent(bool isObjectDeprecated) { + if (isObjectDeprecated) { + AtomicSet(HasDeprecatedObjects, 1); + } + return isObjectDeprecated ? nullptr : SelfRefsAreDeadPtr; + } + + void CreateSelfRefsAreDeadPtr(); + + using IThreadRef = TAutoPtr<IThreadFactory::IThread>; + +private: + const TPQLibSettings Settings; + + std::shared_ptr<grpc::CompletionQueue> CQ; + std::vector<std::thread> GRpcThreads; + + std::shared_ptr<std::atomic<bool>> ShuttingDown = std::make_shared<std::atomic<bool>>(false); // shared_ptr is for deprecated case when we have leaked threads + + THolder<TScheduler> Scheduler; + TQueuePool QueuePool; + std::shared_ptr<IThreadPool> CompressionPool; + + TIntrusivePtr<ILogger> Logger; + + TAtomic HasDeprecatedObjects = 0; + + TAutoEvent SelfRefsAreDeadEvent; + std::shared_ptr<void> SelfRefsAreDeadPtr; + TAdaptiveLock Lock, UALock; + size_t DestroySetAddCounter = 0; + std::set<std::weak_ptr<TSyncDestroyed>, std::owner_less<std::weak_ptr<TSyncDestroyed>>> SyncDestroyedSet; + + TString UserAgent = "C++ pqlib v0.3"; + + struct TYdbWrapper { + THolder<NYdb::NPersQueue::TPersQueueClient> PersQueueClient; + }; + + THolder<NYdb::TDriver> YdbDriver; + std::map<TWrapperKey, TYdbWrapper> Wrappers; + NYdb::NPersQueue::IExecutor::TPtr CompressExecutor; + NYdb::NPersQueue::IExecutor::TPtr HandlersExecutor; + std::shared_ptr<NYdb::ICredentialsProviderFactory> CredentialsProviderFactoryWrapper; +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/persqueue_p_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/persqueue_p_ut.cpp new file mode 100644 index 0000000000..a695c2a8da --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/persqueue_p_ut.cpp @@ -0,0 +1,8 @@ +#include "persqueue_p.h" +#include "producer.h" +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h> +#include <library/cpp/testing/unittest/registar.h> + +namespace NPersQueue { +} +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/processor.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/processor.cpp new file mode 100644 index 0000000000..1b3ca57fd8 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/processor.cpp @@ -0,0 +1,358 @@ +#include "processor.h" +#include <util/string/builder.h> + +namespace NPersQueue { + +template<typename T> +bool IsReady(const NThreading::TFuture<T>& future) +{ + return future.HasValue() || future.HasException(); +} + +template<typename T> +void Reset(NThreading::TFuture<T>& future) +{ + future = NThreading::TFuture<T>(); +} + +template<typename T> +void TProcessor::SafeSubscribe(NThreading::TFuture<T>& future) +{ + std::weak_ptr<TProcessor> self(shared_from_this()); + PQLib->Subscribe( + future, + this, + [self](const auto&) { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->ProcessFutures(); + }}); +} + +TProcessor::TProcessor(const TProcessorSettings& settings, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger) noexcept + : IProcessorImpl(std::move(destroyEventRef), std::move(pqLib)) + , Settings(settings) + , Logger(std::move(logger)) +{} + +TProcessor::~TProcessor() noexcept +{ + SubscribeDestroyed(); + CleanState(); +} + +void TProcessor::Init() +{ + auto selfShared = shared_from_this(); + auto init = [selfShared] { + selfShared->InitImpl(); + }; + Y_VERIFY(PQLib->GetQueuePool().GetQueue(this).AddFunc(init)); +} + +void TProcessor::InitImpl() +{ + ScheduleIdleSessionsCleanup(); + RecreateConsumer(); + SafeSubscribe(ConsumerIsStarted); + SafeSubscribe(ConsumerIsDead); +} + +// Must be called under lock +void TProcessor::CleanState() noexcept +{ + ProcessingData.clear(); + Reset(ConsumerIsDead); + Reset(ConsumerRequest); + Consumer.reset(); + Producers.clear(); + Data.clear(); + CurrentOriginDataMemoryUsage = 0; + CurrentProcessedDataMemoryUsage = 0; +} + + +bool TProcessor::NextRequest() noexcept +{ + Y_VERIFY(Consumer); + Y_VERIFY(!ConsumerRequest.Initialized()); + + DEBUG_LOG("Current memory usage: " + << "origin messages " << CurrentOriginDataMemoryUsage << "(" << Settings.MaxOriginDataMemoryUsage << "), " + << "processed messages " << CurrentProcessedDataMemoryUsage << "(" << Settings.MaxProcessedDataMemoryUsage << ")\n", + "", ""); + + if (CurrentOriginDataMemoryUsage < Settings.MaxOriginDataMemoryUsage && CurrentProcessedDataMemoryUsage < Settings.MaxProcessedDataMemoryUsage) { + DEBUG_LOG("Issuing consumer request\n", "", ""); + ConsumerRequest = Consumer->GetNextMessage(); + return true; + } else { + return false; + } +} + +bool TProcessor::ProcessConsumerResponse(TOriginData& result, TVector<NThreading::TFuture<TProcessedData>>& processedDataFutures) noexcept +{ + Y_VERIFY(ConsumerRequest.HasValue()); + auto message = ConsumerRequest.ExtractValueSync(); + Reset(ConsumerRequest); + const auto& type = message.Type; + const auto& resp = message.Response; + if (type == EMT_ERROR) { + // Will be handled later via ConsumerIsDead + ERR_LOG("Got error: " << resp << "\n", "", ""); + } else if (type == EMT_RELEASE) { + Y_VERIFY(resp.HasRelease()); + auto t = resp.GetRelease().GetTopic(); + auto p = resp.GetRelease().GetPartition(); + auto g = resp.GetRelease().GetGeneration(); + DEBUG_LOG("Got release for " << t << ":" << p << ", generation " << g << "\n", "", ""); + } else if (type == EMT_LOCK) { + Y_VERIFY(resp.HasLock()); + auto t = resp.GetLock().GetTopic(); + auto p = resp.GetLock().GetPartition(); + auto g = resp.GetLock().GetGeneration(); + DEBUG_LOG("Got lock for " << t << ":" << p << ", generation " << g << "\n", "", ""); + message.ReadyToRead.SetValue(TLockInfo{}); + } else if (type == EMT_DATA) { + Y_VERIFY(resp.HasData()); + Data.push_back(TDataInfo{resp.GetData().GetCookie()}); + for (auto& batch: resp.GetData().GetMessageBatch()) { + TPartition partition(batch.GetTopic(), batch.GetPartition()); + TVector<TOriginMessage>& originMessages = result.Messages[partition]; + TDeque<TProcessingData>& processingData = ProcessingData[partition]; + Data.back().TotalOriginMessages += batch.MessageSize(); + + for (auto& msg: batch.GetMessage()) { + CurrentOriginDataMemoryUsage += msg.ByteSizeLong(); + NThreading::TPromise<TProcessedData> promise = NThreading::NewPromise<TProcessedData>(); + NThreading::TFuture<TProcessedData> future = promise.GetFuture(); + TOriginMessage originMessage = TOriginMessage{msg, promise}; + processedDataFutures.push_back(future); + originMessages.push_back(originMessage); + processingData.push_back(TProcessingData{future, msg.GetOffset(), msg.ByteSizeLong(), &Data.back()}); + } + } + DEBUG_LOG("Processing data message " << resp.GetData().GetCookie() << "\n", "", ""); + return true; + } else if (type == EMT_COMMIT) { + DEBUG_LOG("Got commit", "", ""); + } else { + WARN_LOG("Unknown message type " << int(type) << ", response " << resp << "\n", "", ""); + } + + return false; +} + +// Must be called under lock +void TProcessor::RecreateConsumer() noexcept +{ + if (DestroyedFlag) { + return; + } + + INFO_LOG("Create consumer\n", "", ""); + CleanState(); + + TConsumerSettings consumerSettings = Settings.ConsumerSettings; + consumerSettings.UseLockSession = true; + + Consumer = PQLib->CreateConsumer(consumerSettings, DestroyEventRef == nullptr, Logger); + ConsumerIsStarted = Consumer->Start(); + ConsumerIsDead = Consumer->IsDead(); +} + +void TProcessor::ScheduleIdleSessionsCleanup() noexcept +{ + if (DestroyedFlag) { + return; + } + + std::weak_ptr<TProcessor> self(shared_from_this()); + PQLib->GetScheduler().Schedule(Settings.SourceIdIdleTimeout, + this, + [self] { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->CloseIdleSessions(); + } + }); +} + +void TProcessor::CloseIdleSessions() noexcept +{ + if (DestroyedFlag) { + return; + } + + TInstant now = TInstant::Now(); + for (auto it = Producers.begin(); it != Producers.end();) { + if (it->second.LastWriteTime + Settings.SourceIdIdleTimeout < now) { + INFO_LOG("Close producer for sourceid=" << it->first.SourceId, "", ""); + it = Producers.erase(it); + } else { + ++it; + } + } + ScheduleIdleSessionsCleanup(); +} + +void TProcessor::ProcessFutures() noexcept +{ + if (DestroyedFlag) { + return; + } + + TMaybe<NThreading::TFuture<TConsumerMessage>> consumerRequestFuture; + TVector<NThreading::TFuture<TProcessedData>> processedDataFutures; + TVector<NThreading::TFuture<TProducerCommitResponse>> producerAckFutures; + + if (IsReady(ConsumerIsDead)) { + ERR_LOG("Consumer died with error" << ConsumerIsDead.ExtractValueSync(), "", ""); + RecreateConsumer(); + SafeSubscribe(ConsumerIsStarted); + SafeSubscribe(ConsumerIsDead); + return; + } + + if (IsReady(ConsumerRequest)) { + TOriginData data; + bool hasData = ProcessConsumerResponse(data, processedDataFutures); + if (hasData) { + Requests.front().SetValue(data); + Requests.pop_front(); + } + } + + TInstant now = TInstant::Now(); + for (auto& partitionData: ProcessingData) { + TDeque<TProcessingData>& processingData = partitionData.second; + + while (!processingData.empty() && IsReady(processingData.front().Processed)) { + Y_VERIFY(processingData.front().Processed.HasValue(), "processing future cannot be filled with exception"); + ui64 originMessageOffset = processingData.front().OriginMessageOffset; + CurrentOriginDataMemoryUsage -= processingData.front().OriginMessageSize; + TProcessedData processedData = processingData.front().Processed.ExtractValueSync(); + TDataInfo* dataInfo = processingData.front().DataInfo; + dataInfo->OriginMessagesCounter++; + dataInfo->TotalProcessedMessages += processedData.Messages.size(); + processingData.pop_front(); + for (auto& processedMessage: processedData.Messages) { + TString newSourceId = processedMessage.SourceIdPrefix + partitionData.first.Topic + ":" + ToString(partitionData.first.PartitionId) + "_" + processedMessage.Topic; + TProducerKey producerKey { processedMessage.Topic, newSourceId }; + IProducer* producer = nullptr; + + auto it = Producers.find(producerKey); + if (it == Producers.end()) { + TProducerSettings producerSettings = Settings.ProducerSettings; + producerSettings.Topic = processedMessage.Topic; + producerSettings.SourceId = newSourceId; + producerSettings.PartitionGroup = processedMessage.Group; + producerSettings.ReconnectOnFailure = true; + Producers[producerKey] = TProducerInfo { PQLib->CreateRetryingProducer(producerSettings, DestroyEventRef == nullptr, Logger), {}, now}; + producer = Producers[producerKey].Producer.get(); + producer->Start(); + } else { + producer = it->second.Producer.get(); + } + + CurrentProcessedDataMemoryUsage += processedMessage.Data.size(); + DEBUG_LOG("Write message with seqNo=" << originMessageOffset + 1 << " to sourceId=" << newSourceId, "", ""); + NThreading::TFuture<TProducerCommitResponse> ack = producer->Write(originMessageOffset + 1, std::move(processedMessage.Data)); + producerAckFutures.push_back(ack); + Producers[producerKey].Queue.push_back(TProduceInfo(ack, dataInfo)); + } + } + } + + for (auto& producer: Producers) { + TProducerInfo& producerInfo = producer.second; + while (!producerInfo.Queue.empty() && IsReady(producerInfo.Queue.front().Ack)) { + auto ack = producerInfo.Queue.front().Ack.ExtractValueSync(); + DEBUG_LOG("Got ack for message with seqNo=" << ack.SeqNo << " from sourceId=" << producer.first.SourceId << " " << ack.Response, "", ""); + producerInfo.LastWriteTime = now; + producerInfo.Queue.front().DataInfo->ProcessedMessagesCounter++; + CurrentProcessedDataMemoryUsage -= ack.Data.GetSourceData().size(); + producerInfo.Queue.pop_front(); + } + } + + TVector<ui64> toCommit; + while (!Data.empty() && Data.front().IsReadyForCommit()) { + toCommit.push_back(Data.front().Cookie); + Data.pop_front(); + } + if (!toCommit.empty()) { + Consumer->Commit(toCommit); + } + + if (IsReady(ConsumerIsStarted) && !ConsumerRequest.Initialized() && !Requests.empty()) { + if (NextRequest()) { + consumerRequestFuture = ConsumerRequest; + } + } + + if (consumerRequestFuture.Defined()) { + SafeSubscribe(*consumerRequestFuture.Get()); + } + + for (auto& future: processedDataFutures) { + SafeSubscribe(future); + } + + for (auto& future: producerAckFutures) { + SafeSubscribe(future); + } +} + +void TProcessor::GetNextData(NThreading::TPromise<TOriginData>& promise) noexcept +{ + Y_VERIFY(!DestroyedFlag); + TMaybe<NThreading::TFuture<TConsumerMessage>> consumerRequestFuture; + Requests.push_back(promise); + if (IsReady(ConsumerIsStarted) && !ConsumerRequest.Initialized()) { + if (NextRequest()) { + consumerRequestFuture = ConsumerRequest; + } + } + + if (consumerRequestFuture.Defined()) { + SafeSubscribe(*consumerRequestFuture.Get()); + } +} + +void TProcessor::SubscribeDestroyed() { + NThreading::TPromise<void> promise = ObjectsDestroyed; + auto handler = [promise](const auto&) mutable { + promise.SetValue(); + }; + std::vector<NThreading::TFuture<void>> futures; + futures.reserve(Producers.size() + 2); + futures.push_back(DestroyedPromise.GetFuture()); + if (Consumer) { + futures.push_back(Consumer->Destroyed()); + } + for (const auto& p : Producers) { + if (p.second.Producer) { + futures.push_back(p.second.Producer->Destroyed()); + } + } + WaitExceptionOrAll(futures).Subscribe(handler); +} + +NThreading::TFuture<void> TProcessor::Destroyed() noexcept { + return ObjectsDestroyed.GetFuture(); +} + +void TProcessor::Cancel() { + DestroyedFlag = true; + + for (auto& req : Requests) { + req.SetValue(TOriginData()); + } + Requests.clear(); + + DestroyPQLibRef(); +} + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/processor.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/processor.h new file mode 100644 index 0000000000..a6abb827a0 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/processor.h @@ -0,0 +1,118 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include "persqueue_p.h" + +#include <library/cpp/threading/future/future.h> +#include <util/generic/deque.h> +#include <util/generic/map.h> +#include <util/generic/string.h> +#include <util/generic/vector.h> + +namespace NPersQueue { + +struct TDataInfo { + ui64 Cookie; + ui64 TotalOriginMessages; + ui64 OriginMessagesCounter; + ui64 TotalProcessedMessages; + ui64 ProcessedMessagesCounter; + + TDataInfo(ui64 cookie) + : Cookie(cookie) + , TotalOriginMessages(0) + , OriginMessagesCounter(0) + , TotalProcessedMessages(0) + , ProcessedMessagesCounter(0) + {} + + bool IsReadyForCommit() { + return TotalOriginMessages == OriginMessagesCounter && TotalProcessedMessages == ProcessedMessagesCounter; + } +}; + +struct TProducerKey { + TString Topic; + TString SourceId; + + bool operator <(const TProducerKey& other) const { + return std::tie(Topic, SourceId) < std::tie(other.Topic, other.SourceId); + } +}; + +struct TProduceInfo { + NThreading::TFuture<TProducerCommitResponse> Ack; + TDataInfo* DataInfo = nullptr; + + explicit TProduceInfo(const NThreading::TFuture<TProducerCommitResponse>& ack, TDataInfo* dataInfo) + : Ack(ack) + , DataInfo(dataInfo) + {} +}; + +struct TProducerInfo { + std::shared_ptr<IProducerImpl> Producer; + TDeque<TProduceInfo> Queue; + TInstant LastWriteTime; +}; + +struct TProcessingData { + NThreading::TFuture<TProcessedData> Processed; + ui64 OriginMessageOffset; + ui64 OriginMessageSize; + TDataInfo* DataInfo; + + TProcessingData(const NThreading::TFuture<TProcessedData>& processed, ui64 originMessageOffset, ui64 originMessageSize, TDataInfo* dataInfo) + : Processed(processed) + , OriginMessageOffset(originMessageOffset) + , OriginMessageSize(originMessageSize) + , DataInfo(dataInfo) + {}; +}; + +class TPQLibPrivate; + +class TProcessor: public IProcessorImpl, public std::enable_shared_from_this<TProcessor> { +public: + TProcessor(const TProcessorSettings& settings, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger) noexcept; + ~TProcessor() noexcept; + void Init() override; + + using IProcessorImpl::GetNextData; + void GetNextData(NThreading::TPromise<TOriginData>& promise) noexcept override; + + NThreading::TFuture<void> Destroyed() noexcept override; + + void Cancel() override; + +protected: + TProcessorSettings Settings; + TIntrusivePtr<ILogger> Logger; + std::shared_ptr<IConsumerImpl> Consumer; + NThreading::TFuture<TConsumerCreateResponse> ConsumerIsStarted; + NThreading::TFuture<TError> ConsumerIsDead; + NThreading::TFuture<TConsumerMessage> ConsumerRequest; + TMap<TProducerKey, TProducerInfo> Producers; + ui64 CurrentOriginDataMemoryUsage; + ui64 CurrentProcessedDataMemoryUsage; + + TDeque<NThreading::TPromise<TOriginData>> Requests; + TDeque<TDataInfo> Data; + TMap<TPartition, TDeque<TProcessingData>> ProcessingData; + NThreading::TPromise<void> ObjectsDestroyed = NThreading::NewPromise<void>(); + bool DestroyedFlag = false; + + template<typename T> + void SafeSubscribe(NThreading::TFuture<T>& future); + void CleanState() noexcept; + bool NextRequest() noexcept; + void RecreateConsumer() noexcept; + void ScheduleIdleSessionsCleanup() noexcept; + void CloseIdleSessions() noexcept; + void ProcessFutures() noexcept; + bool ProcessConsumerResponse(TOriginData& result, TVector<NThreading::TFuture<TProcessedData>>& processedDataFutures) noexcept; + void SubscribeDestroyed(); + void InitImpl(); +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/processor_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/processor_ut.cpp new file mode 100644 index 0000000000..d36573dcfc --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/processor_ut.cpp @@ -0,0 +1,133 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h> + +#include <library/cpp/testing/unittest/registar.h> + +using namespace NThreading; +using namespace NKikimr; +using namespace NKikimr::NPersQueueTests; + +namespace NPersQueue { +Y_UNIT_TEST_SUITE(TProcessorTest) { + TConsumerSettings MakeConsumerSettings(const TString& clientId, const TVector<TString>& topics, + const TTestServer& testServer) { + TConsumerSettings consumerSettings; + consumerSettings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + consumerSettings.ClientId = clientId; + consumerSettings.Topics = topics; + consumerSettings.UseLockSession = false; + return consumerSettings; + } + + TProducerSettings MakeProducerSettings(const TString& topic, const TTestServer& testServer) { + TProducerSettings producerSettings; + producerSettings.ReconnectOnFailure = true; + producerSettings.Topic = topic; + producerSettings.SourceId = "123"; + producerSettings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + producerSettings.Codec = ECodec::RAW; + return producerSettings; + } + + TProcessorSettings MakeProcessorSettings(const TTestServer& testServer) { + TProcessorSettings processorSettings; + processorSettings.ConsumerSettings = MakeConsumerSettings("processor", {"input"}, testServer); + processorSettings.ProducerSettings = MakeProducerSettings("output", testServer); + return processorSettings; + } + + void AssertWriteValid(const NThreading::TFuture<TProducerCommitResponse>& respFuture) { + const TProducerCommitResponse& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TWriteResponse::kAck, "Msg: " << resp.Response); + } + + void ProcessMessage(TOriginData& originData) { + for (auto& messages: originData.Messages) { + for (auto& message: messages.second) { + auto msg = message.Message; + + TVector<TProcessedMessage> processedMessages; + for (int i = 0; i < 3; ++i) { + TStringBuilder sourceIdPrefix; + sourceIdPrefix << "shard" << i << "_"; + processedMessages.push_back(TProcessedMessage {"output", ToString(msg.GetOffset()), sourceIdPrefix, 0}); + } + message.Processed.SetValue(TProcessedData{processedMessages}); + } + } + } + + bool IsFinished(const TTestServer& testServer) { + auto clientInfo = testServer.AnnoyingClient->GetClientInfo({"rt3.dc1--input"}, "processor", true); + auto topicResult = clientInfo.GetMetaResponse().GetCmdGetReadSessionsInfoResult().GetTopicResult(0).GetPartitionResult(); + for (auto& partition: topicResult) { + if (partition.GetClientOffset() < partition.GetEndOffset()) return false; + } + return true; + } + + void BaseFunctionalityTest() { + TTestServer testServer(false); + testServer.StartServer(false); + + testServer.AnnoyingClient->FullInit(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--input", 2); + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--output", 4); + + auto producerSettings = MakeProducerSettings("input", testServer); + auto producer = testServer.PQLib->CreateProducer(producerSettings, testServer.PQLibSettings.DefaultLogger, false); + producer->Start(); + + for (int i = 1; i < 4; ++i) { + auto write1 = producer->Write(i, TString("blob")); + AssertWriteValid(write1); + } + + auto processorSettings = MakeProcessorSettings(testServer); + auto processor = testServer.PQLib->CreateProcessor(processorSettings, testServer.PQLibSettings.DefaultLogger, false); + + while (!IsFinished(testServer)) { + auto data = processor->GetNextData(); + data.Subscribe([](const auto& f) { + auto value = f.GetValueSync(); + ProcessMessage(value); + }); + Sleep(TDuration::MilliSeconds(100)); + } + + auto consumerSettings = MakeConsumerSettings("checker", {"output"}, testServer); + auto consumer = testServer.PQLib->CreateConsumer(consumerSettings, nullptr, false); + auto isStarted = consumer->Start().ExtractValueSync(); + UNIT_ASSERT_C(isStarted.Response.HasInit(), "cannot start consumer " << isStarted.Response); + + ui32 totalMessages = 9; + ui32 currentMessages = 0; + THashMap<TString, ui32> expectedSeqno = { + {"shard0_rt3.dc1--input:0_output", 1}, {"shard1_rt3.dc1--input:0_output", 1}, {"shard2_rt3.dc1--input:0_output", 1}, + {"shard0_rt3.dc1--input:1_output", 1}, {"shard1_rt3.dc1--input:1_output", 1}, {"shard2_rt3.dc1--input:1_output", 1}, + {"shard0_rt3.dc1--input:2_output", 1}, {"shard1_rt3.dc1--input:2_output", 1}, {"shard2_rt3.dc1--input:2_output", 1}, + {"shard0_rt3.dc1--input:3_output", 1}, {"shard1_rt3.dc1--input:3_output", 1}, {"shard2_rt3.dc1--input:3_output", 1}, + }; + while (currentMessages < totalMessages) { + auto f = consumer->GetNextMessage(); + for (auto& batch: f.GetValueSync().Response.GetData().GetMessageBatch()) { + for (auto& message: batch.GetMessage()) { + TString sourceId = message.GetMeta().GetSourceId(); + ui32 seqNo = message.GetMeta().GetSeqNo(); + UNIT_ASSERT_EQUAL(ToString(seqNo - 1), message.GetData()); + auto it = expectedSeqno.find(sourceId); + UNIT_ASSERT_UNEQUAL_C(it, expectedSeqno.end(), "unknown sourceId " << sourceId); + UNIT_ASSERT_EQUAL(it->second, seqNo); + it->second++; + currentMessages += 1; + } + } + } + } + + Y_UNIT_TEST(BasicFunctionality) { + BaseFunctionalityTest(); + } +} +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/producer.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/producer.cpp new file mode 100644 index 0000000000..2de983453f --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/producer.cpp @@ -0,0 +1,506 @@ +#include <util/generic/strbuf.h> +#include <util/stream/zlib.h> +#include <util/string/cast.h> +#include <util/string/printf.h> +#include <util/string/vector.h> +#include <util/string/builder.h> + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/impl/producer.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/impl/persqueue_p.h> + +#include <grpc++/create_channel.h> + + +namespace NPersQueue { + +class TProducerDestroyHandler : public IHandler { +public: + TProducerDestroyHandler(std::weak_ptr<TProducer> ptr, TIntrusivePtr<TProducer::TRpcStuff> rpcStuff, + const void* queueTag, TPQLibPrivate* pqLib) + : Ptr(std::move(ptr)) + , RpcStuff(std::move(rpcStuff)) + , QueueTag(queueTag) + , PQLib(pqLib) + {} + + void Destroy(const TError&) override { + auto producer = Ptr; + auto handler = [producer] { + auto producerShared = producer.lock(); + if (producerShared) { + producerShared->OnFail(); + } + }; + Y_VERIFY(PQLib->GetQueuePool().GetQueue(QueueTag).AddFunc(handler)); + } + + TString ToString() override { return "TProducerDestroyHandler"; } + +protected: + std::weak_ptr<TProducer> Ptr; + TIntrusivePtr<TProducer::TRpcStuff> RpcStuff; // must simply live + const void* QueueTag; + TPQLibPrivate* PQLib; +}; + +class TStreamCreated : public TProducerDestroyHandler { +public: + TStreamCreated(std::weak_ptr<TProducer> ptr, TIntrusivePtr<TProducer::TRpcStuff> rpcStuff, + const void* queueTag, TPQLibPrivate* pqLib) + : TProducerDestroyHandler(std::move(ptr), std::move(rpcStuff), queueTag, pqLib) + {} + + void Done() override { + auto producer = Ptr; + auto UA = PQLib->GetUserAgent(); + auto handler = [producer, UA] { + auto producerShared = producer.lock(); + if (producerShared) { + producerShared->OnStreamCreated(UA); + } + }; + Y_VERIFY(PQLib->GetQueuePool().GetQueue(QueueTag).AddFunc(handler)); + } + + TString ToString() override { return "StreamCreated"; } +}; + +class TFinishDone : public TProducerDestroyHandler { +public: + TFinishDone(std::weak_ptr<TProducer> ptr, TIntrusivePtr<TProducer::TRpcStuff> rpcStuff, + const void* queueTag, TPQLibPrivate* pqLib) + : TProducerDestroyHandler(std::move(ptr), std::move(rpcStuff), queueTag, pqLib) + {} + + void Done() override { + auto producer = Ptr; + auto handler = [producer] { + auto producerShared = producer.lock(); + if (producerShared) { + producerShared->OnFinishDone(); + } + }; + Y_VERIFY(PQLib->GetQueuePool().GetQueue(QueueTag).AddFunc(handler)); + } + + void Destroy(const TError&) override { + Y_FAIL("Finish call failed"); + } + + TString ToString() override { return "Finish"; } +}; + + +class TWriteDone : public TProducerDestroyHandler { +public: + TWriteDone(std::weak_ptr<TProducer> ptr, TIntrusivePtr<TProducer::TRpcStuff> rpcStuff, + const void* queueTag, TPQLibPrivate* pqLib) + : TProducerDestroyHandler(std::move(ptr), std::move(rpcStuff), queueTag, pqLib) + {} + + void Done() override { + auto producer = Ptr; + auto handler = [producer] { + auto producerShared = producer.lock(); + if (producerShared) { + producerShared->OnWriteDone(); + } + }; + Y_VERIFY(PQLib->GetQueuePool().GetQueue(QueueTag).AddFunc(handler)); + } + TString ToString() override { return "WriteDone"; } +}; + + +class TReadDone : public TProducerDestroyHandler { +public: + TReadDone(std::weak_ptr<TProducer> ptr, TIntrusivePtr<TProducer::TRpcStuff> rpcStuff, + const void* queueTag, TPQLibPrivate* pqLib) + : TProducerDestroyHandler(std::move(ptr), std::move(rpcStuff), queueTag, pqLib) + {} + + void Done() override { + auto producer = Ptr; + auto handler = [producer] { + auto producerShared = producer.lock(); + if (producerShared) { + producerShared->OnReadDone(); + } + }; + Y_VERIFY(PQLib->GetQueuePool().GetQueue(QueueTag).AddFunc(handler)); + } + TString ToString() override { return "ReadDone"; } +}; + +TProducer::TProducer(const TProducerSettings& settings, std::shared_ptr<grpc::CompletionQueue> cq, + NThreading::TPromise<TProducerCreateResponse> promise, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger) noexcept + : IProducerImpl(std::move(destroyEventRef), std::move(pqLib)) + , RpcStuff(new TRpcStuff()) + , Settings(settings) + , StartPromise(promise) + , IsDeadPromise(NThreading::NewPromise<TError>()) + , Pos(0) + , WriteInflight(false) + , ProxyCookie(0) + , Logger(std::move(logger)) + , IsDestroyed(false) + , IsDestroying(false) + , Failing(false) +{ + RpcStuff->CQ = std::move(cq); +} + +TProducer::~TProducer() { + Destroy(); +} + +void TProducer::Destroy() noexcept { + IsDestroying = true; + RpcStuff->Context.TryCancel(); + ChannelHolder.ChannelPtr = nullptr; + + Destroy("Destructor called"); +} + +void TProducer::DoStart(TInstant deadline) { + if (IsDestroyed) { + return; + } + + if (deadline != TInstant::Max()) { + std::weak_ptr<TProducer> self = shared_from_this(); + auto onStartTimeout = [self] { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->OnStartTimeout(); + } + }; + StartDeadlineCallback = + PQLib->GetScheduler().Schedule(deadline, this, onStartTimeout); + } + + FillMetaHeaders(RpcStuff->Context, Settings.Server.Database, Settings.CredentialsProvider.get()); + RpcStuff->Stub = PersQueueService::NewStub(RpcStuff->Channel); + RpcStuff->Stream = RpcStuff->Stub->AsyncWriteSession(&RpcStuff->Context, RpcStuff->CQ.get(), new TQueueEvent(StreamCreatedHandler)); +} + +void TProducer::Init() { + std::weak_ptr<TProducer> self = shared_from_this(); // we can't make this object in constructor, because this will be the only reference to us and ~shared_ptr() will destroy us. + StreamCreatedHandler.Reset(new TStreamCreated(self, RpcStuff, this, PQLib.Get())); + ReadDoneHandler.Reset(new TReadDone(self, RpcStuff, this, PQLib.Get())); + WriteDoneHandler.Reset(new TWriteDone(self, RpcStuff, this, PQLib.Get())); + FinishDoneHandler.Reset(new TFinishDone(self, RpcStuff, this, PQLib.Get())); +} + +NThreading::TFuture<TProducerCreateResponse> TProducer::Start(TInstant deadline) noexcept { + if (IsDestroyed) { + TWriteResponse res; + res.MutableError()->SetDescription("Producer is dead"); + res.MutableError()->SetCode(NErrorCode::ERROR); + return NThreading::MakeFuture<TProducerCreateResponse>(TProducerCreateResponse{std::move(res)}); + } + NThreading::TFuture<TProducerCreateResponse> ret = StartPromise.GetFuture(); + if (ChannelHolder.ChannelInfo.Initialized()) { + std::weak_ptr<TProducer> self = shared_from_this(); + PQLib->Subscribe(ChannelHolder.ChannelInfo, + this, + [self, deadline](const auto& f) { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->SetChannel(f.GetValue()); + selfShared->DoStart(deadline); + } + }); + } else { + if (RpcStuff->Channel && ProxyCookie) { + DoStart(deadline); + } else { + Destroy("No channel"); + } + } + return ret; +} + +void TProducer::SetChannel(const TChannelHolder& channel) noexcept { + if (IsDestroyed) { + return; + } + + ChannelHolder = channel; +} + +void TProducer::SetChannel(const TChannelInfo& channel) noexcept { + if (IsDestroyed) { + return; + } + + Y_VERIFY(!RpcStuff->Channel); + if (!channel.Channel || channel.ProxyCookie == 0) { + if (channel.Channel) { + Destroy("ProxyCookie is zero"); + } else { + Destroy("Channel creation error"); + } + + return; + } + + RpcStuff->Channel = channel.Channel; + ProxyCookie = channel.ProxyCookie; +} + + +static TProducerCommitResponse MakeErrorCommitResponse(const TString& description, TProducerSeqNo seqNo, const TData& data) { + TWriteResponse res; + res.MutableError()->SetDescription(description); + res.MutableError()->SetCode(NErrorCode::ERROR); + return TProducerCommitResponse{seqNo, data, std::move(res)}; +} + +void TProducer::Write(NThreading::TPromise<TProducerCommitResponse>& promise, TData data) noexcept { + return TProducer::Write(promise, MaxSeqNo + 1, std::move(data)); +} + +void TProducer::Write(NThreading::TPromise<TProducerCommitResponse>& promise, const TProducerSeqNo seqNo, TData data) noexcept { + Y_VERIFY(data.IsEncoded()); + + MaxSeqNo = Max(seqNo, MaxSeqNo); + + if (IsDestroyed) { + promise.SetValue(MakeErrorCommitResponse("producer is dead", seqNo, data)); + return; + } + + if (StartPromise.Initialized()) { + promise.SetValue(MakeErrorCommitResponse("producer is not ready", seqNo, data)); + return; + } + + CommitPromises.push_back(promise); + Data.emplace_back(seqNo, std::move(data)); + DoWrite(); +} + +void TProducer::OnWriteDone() { + if (IsDestroyed) { + return; + } + + WriteInflight = false; + DoWrite(); +} + +void TProducer::DoWrite() { + if (Failing && !WriteInflight) { + return; + } + if (WriteInflight || Pos == Data.size()) { + return; + } + WriteInflight = true; + TWriteRequest req; + Settings.CredentialsProvider->FillAuthInfo(req.MutableCredentials()); + + req.MutableData()->SetSeqNo(Data[Pos].SeqNo); + req.MutableData()->SetData(Data[Pos].Data.GetEncodedData()); + req.MutableData()->SetCodec(Data[Pos].Data.GetCodecType()); + req.MutableData()->SetCreateTimeMs(Data[Pos].Data.GetTimestamp().MilliSeconds()); + + ++Pos; + RpcStuff->Stream->Write(req, new TQueueEvent(WriteDoneHandler)); +} + +NThreading::TFuture<TError> TProducer::IsDead() noexcept { + return IsDeadPromise.GetFuture(); +} + +void TProducer::Destroy(const TString& description) { + TError error; + error.SetDescription(description); + error.SetCode(NErrorCode::ERROR); + Destroy(error); +} + +void TProducer::Destroy(const TError& error) { + if (IsDestroyed) { + return; + } + + if (!IsDestroying) { + INFO_LOG("Error: " << error, Settings.SourceId, SessionId); + } + + IsDestroyed = true; + IsDeadPromise.SetValue(error); + + if (StartDeadlineCallback) { + StartDeadlineCallback->TryCancel(); + } + + NThreading::TFuture<TChannelInfo> tmp; + tmp.Swap(ChannelHolder.ChannelInfo); + ChannelHolder.ChannelPtr = nullptr; + + Error = error; + if (StartPromise.Initialized()) { + NThreading::TPromise<TProducerCreateResponse> tmp; + tmp.Swap(StartPromise); + TWriteResponse res; + res.MutableError()->CopyFrom(Error); + tmp.SetValue(TProducerCreateResponse{std::move(res)}); + } + while (!CommitPromises.empty()) { + auto p = CommitPromises.front(); + CommitPromises.pop_front(); + auto pp(std::move(Data.front())); + Data.pop_front(); + if (Pos > 0) { + --Pos; + } + TWriteResponse res; + res.MutableError()->CopyFrom(Error); + + p.SetValue(TProducerCommitResponse{pp.SeqNo, std::move(pp.Data), std::move(res)}); + } + + StreamCreatedHandler.Reset(); + ReadDoneHandler.Reset(); + WriteDoneHandler.Reset(); + FinishDoneHandler.Reset(); + + DestroyPQLibRef(); +} + +void TProducer::OnStreamCreated(const TString& userAgent) { + if (IsDestroyed) { + return; + } + + TWriteRequest req; + Settings.CredentialsProvider->FillAuthInfo(req.MutableCredentials()); + + req.MutableInit()->SetSourceId(Settings.SourceId); + req.MutableInit()->SetTopic(Settings.Topic); + req.MutableInit()->SetProxyCookie(ProxyCookie); + req.MutableInit()->SetPartitionGroup(Settings.PartitionGroup); + req.MutableInit()->SetVersion(userAgent); + Y_VERIFY(!userAgent.empty()); + + for (const auto& attr : Settings.ExtraAttrs) { + auto item = req.MutableInit()->MutableExtraFields()->AddItems(); + item->SetKey(attr.first); + item->SetValue(attr.second); + } + + WriteInflight = true; + RpcStuff->Stream->Write(req, new TQueueEvent(WriteDoneHandler)); + + RpcStuff->Response.Clear(); + RpcStuff->Stream->Read(&RpcStuff->Response, new TQueueEvent(ReadDoneHandler)); +} + +void TProducer::OnFail() { + if (Failing) + return; + Failing = true; + + if (IsDestroyed) { + return; + } + + RpcStuff->Stream->Finish(&RpcStuff->Status, new TQueueEvent(FinishDoneHandler)); +} + +void TProducer::OnFinishDone() { + if (IsDestroyed) { + return; + } + + TError error; + const auto& msg = RpcStuff->Status.error_message(); + TString reason(msg.data(), msg.length()); + error.SetDescription(reason); + error.SetCode(NErrorCode::ERROR); + + Destroy(error); +} + +void TProducer::OnReadDone() { + if (IsDestroyed) { + return; + } + + if (RpcStuff->Response.HasError()) { + Destroy(RpcStuff->Response.GetError()); + return; + } + + if (StartPromise.Initialized()) { //init response + NThreading::TPromise<TProducerCreateResponse> tmp; + tmp.Swap(StartPromise); + Y_VERIFY(RpcStuff->Response.HasInit()); + auto res(std::move(RpcStuff->Response)); + RpcStuff->Response.Clear(); + + SessionId = res.GetInit().GetSessionId(); + const ui32 partition = res.GetInit().GetPartition(); + MaxSeqNo = res.GetInit().GetMaxSeqNo(); + const TProducerSeqNo seqNo = MaxSeqNo; + const TString topic = res.GetInit().GetTopic(); + tmp.SetValue(TProducerCreateResponse{std::move(res)}); + + if (StartDeadlineCallback) { + StartDeadlineCallback->TryCancel(); + } + StartDeadlineCallback = nullptr; + + RpcStuff->Stream->Read(&RpcStuff->Response, new TQueueEvent(ReadDoneHandler)); + + DEBUG_LOG("Stream created to topic " << topic << " partition " << partition << " maxSeqNo " << seqNo, Settings.SourceId, SessionId); + } else { //write response + //get first CommitPromise + Y_VERIFY(!Data.empty()); + auto p = CommitPromises.front(); + CommitPromises.pop_front(); + auto pp(std::move(Data.front())); + Data.pop_front(); + Y_VERIFY(Pos > 0); + --Pos; + + Y_VERIFY(RpcStuff->Response.HasAck()); + Y_VERIFY(RpcStuff->Response.GetAck().GetSeqNo() == pp.SeqNo); + auto res(std::move(RpcStuff->Response)); + RpcStuff->Response.Clear(); + RpcStuff->Stream->Read(&RpcStuff->Response, new TQueueEvent(ReadDoneHandler)); + + p.SetValue(TProducerCommitResponse{pp.SeqNo, std::move(pp.Data), std::move(res)}); + } +} + +void TProducer::OnStartTimeout() { + if (IsDestroyed) { + return; + } + + StartDeadlineCallback = nullptr; + if (!StartPromise.Initialized()) { + // everything is OK, there is no timeout, we have already started. + return; + } + + TError error; + error.SetDescription("Start timeout"); + error.SetCode(NErrorCode::CREATE_TIMEOUT); + Destroy(error); +} + +void TProducer::Cancel() { + IsDestroying = true; + RpcStuff->Context.TryCancel(); + ChannelHolder.ChannelPtr = nullptr; + + Destroy(GetCancelReason()); +} + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/producer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/producer.h new file mode 100644 index 0000000000..af74ced4b9 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/producer.h @@ -0,0 +1,107 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/impl/channel.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/impl/iproducer_p.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/impl/internals.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/impl/scheduler.h> +#include <library/cpp/threading/future/future.h> +#include <kikimr/yndx/api/grpc/persqueue.grpc.pb.h> +#include <deque> + +namespace NPersQueue { + +class TProducer : public IProducerImpl, public std::enable_shared_from_this<TProducer> { +public: + using IProducerImpl::Write; + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, const TProducerSeqNo seqNo, TData data) noexcept override; + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TData data) noexcept override; + + NThreading::TFuture<TError> IsDead() noexcept override; + + NThreading::TFuture<TProducerCreateResponse> Start(TInstant deadline) noexcept override; + + TProducer(const TProducerSettings& settings, std::shared_ptr<grpc::CompletionQueue> cq, + NThreading::TPromise<TProducerCreateResponse> promise, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger) noexcept; + ~TProducer(); + + void Destroy() noexcept; + void SetChannel(const TChannelHolder& channel) noexcept; + void SetChannel(const TChannelInfo& channel) noexcept; + + void DoStart(TInstant deadline); + + void Init() override; + + const TProducerSettings& GetSettings() const { + return Settings; + } + + void Cancel() override; + +public: + using TStream = grpc::ClientAsyncReaderWriterInterface<TWriteRequest, TWriteResponse>; + + // objects that must live after destruction of producer untill all the callbacks arrive at CompletionQueue + struct TRpcStuff: public TAtomicRefCount<TRpcStuff> { + TWriteResponse Response; + std::shared_ptr<grpc::CompletionQueue> CQ; + std::shared_ptr<grpc::Channel> Channel; + std::unique_ptr<PersQueueService::Stub> Stub; + grpc::ClientContext Context; + std::unique_ptr<TStream> Stream; + grpc::Status Status; + }; + +private: + friend class TPQLibPrivate; + friend class TStreamCreated; + friend class TReadDone; + friend class TWriteDone; + friend class TProducerDestroyHandler; + friend class TFinishDone; + + IHandlerPtr StreamCreatedHandler; + IHandlerPtr ReadDoneHandler; + IHandlerPtr WriteDoneHandler; + IHandlerPtr FinishDoneHandler; + + void Destroy(const TError& error); + void Destroy(const TString& description); // the same but with Code=ERROR + void OnStreamCreated(const TString& userAgent); + void OnReadDone(); + void OnWriteDone(); + void DoWrite(); + void OnFail(); + void OnFinishDone(); + void OnStartTimeout(); + +protected: + TIntrusivePtr<TRpcStuff> RpcStuff; + + TChannelHolder ChannelHolder; + TProducerSettings Settings; + + TString SessionId; + + NThreading::TPromise<TProducerCreateResponse> StartPromise; + NThreading::TPromise<TError> IsDeadPromise; + std::deque<NThreading::TPromise<TProducerCommitResponse>> CommitPromises; + std::deque<TWriteData> Data; + ui32 Pos = 0; + bool WriteInflight; + ui64 ProxyCookie = 0; + TProducerSeqNo MaxSeqNo = 0; + + TIntrusivePtr<ILogger> Logger; + + TError Error; + + bool IsDestroyed; + bool IsDestroying; + bool Failing; + TIntrusivePtr<TScheduler::TCallbackHandler> StartDeadlineCallback; +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/producer_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/producer_ut.cpp new file mode 100644 index 0000000000..5f3deb880d --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/producer_ut.cpp @@ -0,0 +1,202 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/sdk_test_setup.h> +#include <library/cpp/testing/unittest/registar.h> + +namespace NPersQueue { + +Y_UNIT_TEST_SUITE(TProducerTest) { + Y_UNIT_TEST(NotStartedProducerCanBeDestructed) { + // Test that producer doesn't hang on till shutdown + TPQLib lib; + TProducerSettings settings; + settings.Server = TServerSetting{"localhost"}; + settings.Topic = "topic"; + settings.SourceId = "src"; + lib.CreateProducer(settings, {}, false); + } + + TProducerSettings MakeProducerSettings(const TTestServer& testServer) { + TProducerSettings producerSettings; + producerSettings.ReconnectOnFailure = false; + producerSettings.Topic = "topic1"; + producerSettings.SourceId = "123"; + producerSettings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + producerSettings.Codec = ECodec::LZOP; + return producerSettings; + } + + Y_UNIT_TEST(CancelsOperationsAfterPQLibDeath) { + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(false); + + const size_t partitions = 1; + testServer.AnnoyingClient->FullInit(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", partitions); + + testServer.WaitInit("topic1"); + + auto producer = testServer.PQLib->CreateProducer(MakeProducerSettings(testServer), testServer.PQLibSettings.DefaultLogger, false); + UNIT_ASSERT(!producer->Start().GetValueSync().Response.HasError()); + auto isDead = producer->IsDead(); + UNIT_ASSERT(!isDead.HasValue()); + + auto write1 = producer->Write(1, TString("blob1")); + auto write2 = producer->Write(2, TString("blob2")); + auto write3 = producer->Write(3, TString("blob3")); + auto write4 = producer->Write(4, TString("blob4")); + auto write5 = producer->Write(5, TString("blob5")); + + testServer.PQLib = nullptr; + Cerr << "PQLib destroyed" << Endl; + + UNIT_ASSERT(write1.HasValue()); + UNIT_ASSERT(write2.HasValue()); + UNIT_ASSERT(write3.HasValue()); + UNIT_ASSERT(write4.HasValue()); + UNIT_ASSERT(write5.HasValue()); + + UNIT_ASSERT(write1.GetValue().Response.HasError()); + UNIT_ASSERT(write2.GetValue().Response.HasError()); + UNIT_ASSERT(write3.GetValue().Response.HasError()); + UNIT_ASSERT(write4.GetValue().Response.HasError()); + UNIT_ASSERT(write5.GetValue().Response.HasError()); + + auto write6 = producer->Write(6, TString("blob6")); + UNIT_ASSERT(write6.HasValue()); + UNIT_ASSERT(write6.GetValue().Response.HasError()); + } + + Y_UNIT_TEST(WriteToDeadProducer) { + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(false); + + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLibSettings pqLibSettings; + pqLibSettings.DefaultLogger = logger; + THolder<TPQLib> PQLib = MakeHolder<TPQLib>(pqLibSettings); + + auto producer = PQLib->CreateProducer(MakeProducerSettings(testServer), logger, false); + auto f = producer->Start(); + UNIT_ASSERT(f.GetValueSync().Response.HasError()); + Cerr << f.GetValueSync().Response << "\n"; + auto isDead = producer->IsDead(); + isDead.Wait(); + UNIT_ASSERT(isDead.HasValue()); + + auto write = producer->Write(1, TString("blob")); + + Cerr << write.GetValueSync().Response << "\n"; + UNIT_ASSERT(write.GetValueSync().Response.HasError()); + UNIT_ASSERT_STRINGS_EQUAL(write.GetValueSync().Response.GetError().GetDescription(), "Destroyed"); + } + + Y_UNIT_TEST(Auth_StartProducerWithInvalidTokenFromGrpcMetadataPointOfView_StartFailedAndProduerDiedWithBadRequestErrors) { + TPQLibSettings pqLibSettings; + if (!std::getenv("PERSQUEUE_GRPC_API_V1_ENABLED")) + return; + + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(); + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + pqLibSettings.DefaultLogger = logger; + THolder<TPQLib> PQLib = MakeHolder<TPQLib>(pqLibSettings); + + auto fillCredentialsWithInvalidToken = [](NPersQueue::TCredentials* authInfo) { + authInfo->set_tvm_service_ticket("token\n"); + }; + auto producerSettings = MakeProducerSettings(testServer); + producerSettings.CredentialsProvider = std::make_shared<TCallbackCredentialsProvider>(std::move(fillCredentialsWithInvalidToken)); + auto producer = PQLib->CreateProducer(producerSettings, logger, false); + + + Cerr << "Wait for producer start procedure end" << Endl; + auto startResponse = producer->Start().GetValueSync().Response; + + UNIT_ASSERT(startResponse.has_error()); + UNIT_ASSERT_EQUAL_C(NPersQueue::NErrorCode::ERROR, startResponse.error().code(), startResponse); + + Cerr << "Wait for producer death" << Endl; + auto deathCause = producer->IsDead().GetValueSync(); + + + UNIT_ASSERT_EQUAL_C(NPersQueue::NErrorCode::ERROR, deathCause.code(), deathCause); + } + + Y_UNIT_TEST(Codecs_WriteWithNonDefaultCodecThatRequiresAdditionalConfiguration_ConsumerDiesWithBadRequestError) { + SDKTestSetup setup{"Codecs_WriteWithNonDefaultCodecThatRequiresAdditionalConfiguration_ConsumerIsDeadWithBadRequestError"}; + auto log = setup.GetLog(); + auto producerSettings = setup.GetProducerSettings(); + // TTestServer::AnnoyingClient creates topic with default codecs set: raw, gzip, lzop. zstd not included + producerSettings.Codec = ECodec::ZSTD; + auto producer = setup.GetPQLib()->CreateProducer(producerSettings, nullptr, false); + log << TLOG_INFO << "Wait for producer start"; + auto startResponse = producer->Start().GetValueSync().Response; + UNIT_ASSERT_C(!startResponse.HasError(), startResponse); + + log << TLOG_INFO << "Wait for write response"; + auto writeResponse = producer->Write(NUnitTest::RandomString(250 * 1024, std::rand())).GetValueSync().Response; + + + UNIT_ASSERT(writeResponse.HasError()); + Cerr << writeResponse << "\n"; + UNIT_ASSERT_EQUAL_C(NPersQueue::NErrorCode::BAD_REQUEST, writeResponse.GetError().GetCode(), writeResponse); + } + + void StartProducerWithDiscovery(bool cds, const TString& preferredCluster = {}, bool addBrokenDatacenter = !GrpcV1EnabledByDefault()) { + SDKTestSetup setup{"StartProducerWithDiscovery", false}; + setup.Start(true, addBrokenDatacenter); + auto log = setup.GetLog(); + auto producerSettings = setup.GetProducerSettings(); + if (preferredCluster) { + producerSettings.PreferredCluster = preferredCluster; + } + auto producer = setup.GetPQLib()->CreateProducer(producerSettings, nullptr, false); + auto startResponse = producer->Start().GetValueSync().Response; + UNIT_ASSERT_C(!startResponse.HasError(), startResponse); + + log << TLOG_INFO << "Wait for write response"; + auto writeResponse = producer->Write(1, TString("blob1")).GetValueSync().Response; + UNIT_ASSERT(!writeResponse.HasError()); + + producerSettings.Server.UseLogbrokerCDS = cds ? EClusterDiscoveryUsageMode::Use : EClusterDiscoveryUsageMode::DontUse; + producer = setup.GetPQLib()->CreateProducer(producerSettings, nullptr, false); + startResponse = producer->Start().GetValueSync().Response; + UNIT_ASSERT_C(!startResponse.HasError(), startResponse); + + log << TLOG_INFO << "Wait for write response"; + writeResponse = producer->Write(2, TString("blob2")).GetValueSync().Response; + UNIT_ASSERT(!writeResponse.HasError()); + } + + Y_UNIT_TEST(StartProducerWithCDS) { + StartProducerWithDiscovery(true); + } + + Y_UNIT_TEST(StartProducerWithoutCDS) { + StartProducerWithDiscovery(false); + } + + Y_UNIT_TEST(StartProducerWithCDSAndPreferAvailableCluster) { + StartProducerWithDiscovery(true, "dc1"); + } + + Y_UNIT_TEST(StartProducerWithCDSAndPreferUnavailableCluster) { + StartProducerWithDiscovery(true, "dc2", true); + } + + Y_UNIT_TEST(StartProducerWithCDSAndPreferUnknownCluster) { + SDKTestSetup setup{"PreferUnknownCluster"}; + auto producerSettings = setup.GetProducerSettings(); + producerSettings.Server.UseLogbrokerCDS = EClusterDiscoveryUsageMode::Use; + producerSettings.PreferredCluster = "blablabla"; + auto producer = setup.GetPQLib()->CreateProducer(producerSettings, nullptr, false); + auto startResponse = producer->Start().GetValueSync().Response; + UNIT_ASSERT_C(startResponse.HasError(), startResponse); + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/queue_pool.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/queue_pool.cpp new file mode 100644 index 0000000000..6e1c3f9b2e --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/queue_pool.cpp @@ -0,0 +1,31 @@ +#include "queue_pool.h" + +#include <util/digest/city.h> + +namespace NPersQueue { +void TQueuePool::Start(size_t queuesCount) { + Y_VERIFY(queuesCount > 0); + Queues.resize(queuesCount); + for (auto& queuePtr : Queues) { + queuePtr = std::make_shared<TThreadPool>(); + queuePtr->Start(1); // start one thread for each tag + } +} + +void TQueuePool::Stop() { + for (auto& queuePtr : Queues) { + queuePtr->Stop(); + } + Queues.clear(); +} + +const std::shared_ptr<IThreadPool>& TQueuePool::GetQueuePtr(const void* tag) { + Y_VERIFY(!Queues.empty()); + const size_t queue = static_cast<size_t>(CityHash64(reinterpret_cast<const char*>(&tag), sizeof(tag))) % Queues.size(); + return Queues[queue]; +} + +IThreadPool& TQueuePool::GetQueue(const void* tag) { + return *GetQueuePtr(tag); +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/queue_pool.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/queue_pool.h new file mode 100644 index 0000000000..961b8ff446 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/queue_pool.h @@ -0,0 +1,23 @@ +#pragma once +#include <util/generic/ptr.h> +#include <util/thread/pool.h> + +#include <functional> +#include <memory> +#include <vector> + +namespace NPersQueue { +class TQueuePool { +public: + void Start(size_t queuesCount); + void Stop(); + + // get one-thread-queue for processing specified tag (address) + IThreadPool& GetQueue(const void* tag); + + const std::shared_ptr<IThreadPool>& GetQueuePtr(const void* tag); + +private: + std::vector<std::shared_ptr<IThreadPool>> Queues; +}; +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/queue_pool_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/queue_pool_ut.cpp new file mode 100644 index 0000000000..9f8d792e52 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/queue_pool_ut.cpp @@ -0,0 +1,26 @@ +#include "queue_pool.h" + +#include <library/cpp/testing/unittest/registar.h> + +#include <unordered_map> + +namespace NPersQueue { +Y_UNIT_TEST_SUITE(TQueuePoolTest) { + Y_UNIT_TEST(QueuesDistribution) { + TQueuePool pool; + pool.Start(10); + size_t addresses[10000] = {}; + std::unordered_map<IThreadPool*, size_t> queueAddresses; + for (size_t& address : addresses) { + IThreadPool* q = &pool.GetQueue(&address); + ++queueAddresses[q]; + UNIT_ASSERT_EQUAL(q, &pool.GetQueue(&address)); // one address always leads to one queue + } + UNIT_ASSERT_VALUES_EQUAL(queueAddresses.size(), 10); + for (const auto& queueToCount : queueAddresses) { + UNIT_ASSERT_C(queueToCount.second >= 850, "Count: " << queueToCount.second); + UNIT_ASSERT_C(queueToCount.second <= 1150, "Count: " << queueToCount.second); + } + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_consumer.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_consumer.cpp new file mode 100644 index 0000000000..02b3aafafc --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_consumer.cpp @@ -0,0 +1,526 @@ +#include "retrying_consumer.h" + +namespace { + bool IsRetryable(const NPersQueue::TError& err) { + switch (err.code()) { + case NPersQueue::NErrorCode::INITIALIZING: + case NPersQueue::NErrorCode::OVERLOAD: + case NPersQueue::NErrorCode::READ_TIMEOUT: + case NPersQueue::NErrorCode::TABLET_IS_DROPPED: + case NPersQueue::NErrorCode::CREATE_TIMEOUT: + case NPersQueue::NErrorCode::ERROR: + case NPersQueue::NErrorCode::CLUSTER_DISABLED: + return true; + default: + return false; + } + } +} + +namespace NPersQueue { + +TRetryingConsumer::TRetryingConsumer(const TConsumerSettings& settings, + std::shared_ptr<void> destroyEventRef, + TIntrusivePtr<TPQLibPrivate> pqLib, + TIntrusivePtr<ILogger> logger) + : IConsumerImpl(std::move(destroyEventRef), std::move(pqLib)) + , Settings(settings) + , Logger(std::move(logger)) + , IsDeadPromise(NThreading::NewPromise<TError>()) + , ConsumerDestroyedPromise(NThreading::NewPromise<void>()) + , GenCounter(1) + , CookieCounter(1) + , ReconnectionAttemptsDone(0) + , Stopping(false) + , Reconnecting(false) +{ + Y_VERIFY(Settings.ReconnectOnFailure, "ReconnectOnFailure should be set."); + Y_ENSURE(Settings.MaxAttempts != 0, "MaxAttempts setting can't be zero."); + Y_ENSURE(Settings.ReconnectionDelay != TDuration::Zero(), "ReconnectionDelay setting can't be zero."); + Y_ENSURE(Settings.MaxReconnectionDelay != TDuration::Zero(), "MaxReconnectionDelay setting can't be zero."); + Y_ENSURE(Settings.StartSessionTimeout != TDuration::Zero(), "StartSessionTimeout setting can't be zero."); + + Settings.ReconnectOnFailure = false; // reset flag, to use the settings for creating sub consumer +} + +TRetryingConsumer::~TRetryingConsumer() noexcept { + Cancel(); +} + +NThreading::TFuture<TConsumerCreateResponse> TRetryingConsumer::Start(TInstant deadline) noexcept { + StartPromise = NThreading::NewPromise<TConsumerCreateResponse>(); + if (deadline != TInstant::Max()) { + std::weak_ptr<TRetryingConsumer> weakRef = shared_from_this(); + auto onDeadline = [weakRef] { + auto self = weakRef.lock(); + if (nullptr != self) { + self->OnStartDeadline(); + } + }; + PQLib->GetScheduler().Schedule(deadline, this, onDeadline); + } + DoReconnect(deadline); + return StartPromise.GetFuture(); +} + +NThreading::TFuture<TError> TRetryingConsumer::IsDead() noexcept { + return IsDeadPromise.GetFuture(); +} + +void TRetryingConsumer::GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept { + if (!StartPromise.Initialized() || Stopping) { + TReadResponse res; + res.MutableError()->SetDescription("consumer is not ready"); + res.MutableError()->SetCode(NErrorCode::ERROR); + promise.SetValue(TConsumerMessage{std::move(res)}); + return; + } + + if (!ReadyResponses.empty()) { + Y_ASSERT(PendingRequests.empty()); + promise.SetValue(std::move(ReadyResponses.front())); + ReadyResponses.pop_front(); + return; + } + + PendingRequests.push_back(promise); + + if (Consumer && StartFuture.HasValue()) { + DoRequest(); + } +} + +static void FormatCookies(TStringBuilder& ret, const TVector<ui64>& cookies) { + ret << "{"; + for (size_t i = 0; i < cookies.size(); ++i) { + if (i > 0) { + ret << ", "; + } + ret << cookies[i]; + } + ret << "}"; +} + +static TString FormatCommitForLog(const TStringBuf reason, const TVector<ui64>& cookies, const TVector<ui64>& originalCookies, const TVector<ui64>& committedCookies) { + TStringBuilder ret; + ret << "Commit cookies by retrying consumer" << reason << ". User cookies: "; + FormatCookies(ret, cookies); + ret << ". Subconsumer cookies to commit: "; + FormatCookies(ret, originalCookies); + ret << ". Skipped cookies: "; + FormatCookies(ret, committedCookies); + return std::move(ret); +} + +void TRetryingConsumer::Commit(const TVector<ui64>& cookies) noexcept { + if (!Consumer) { + if (!StartPromise.Initialized() || Stopping) { + Destroy("Not ready", NErrorCode::BAD_REQUEST); + } else { + // just response that cookies were commited + FastCommit(cookies); + DEBUG_LOG(FormatCommitForLog(". Consumer is not initialied", cookies, {}, cookies), "", SessionId); + } + return; + } + if (Settings.CommitsDisabled) { + Destroy("Commits are disabled", NErrorCode::BAD_REQUEST); + return; + } + + TVector<ui64> originalCookies(Reserve(cookies.size())); + // cookies which can be treated as committed + TVector<ui64> commitedCookies; + ui64 minCookie = CookieCounter - Cookies.size(); + for (auto cookie : cookies) { + if (cookie >= minCookie && cookie < CookieCounter) { + auto& cookieInfo = Cookies[cookie - minCookie]; + Y_VERIFY(cookieInfo.UserCookie == cookie); + if (0 != cookieInfo.OriginalCookie) { + CommittingCookies[cookieInfo.OriginalCookie] = cookie; + originalCookies.push_back(cookieInfo.OriginalCookie); + cookieInfo.OriginalCookie = 0; + for (auto* lock : cookieInfo.Locks) { + lock->Cookies.erase(cookie); + } + } else { + commitedCookies.push_back(cookie); + } + } else if (cookie >= CookieCounter) { + Destroy("Unknown cookie", NErrorCode::BAD_REQUEST); + break; + } else { + commitedCookies.push_back(cookie); + } + } + if (!originalCookies.empty()) { + Consumer->Commit(originalCookies); + } + if (!commitedCookies.empty()) { + FastCommit(commitedCookies); + } + DEBUG_LOG(FormatCommitForLog("", cookies, originalCookies, commitedCookies), "", SessionId); + // clean commited cookies + while (!Cookies.empty() && 0 == Cookies.front().OriginalCookie) { + Cookies.pop_front(); + } +} + +void TRetryingConsumer::RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept { + if (!Consumer) { + // client should receive message about release lock + return; + } + auto lockIt = Locks.find(std::make_pair(topic, partition)); + if (lockIt == Locks.end() || !lockIt->second.Locked) { + WARN_LOG("Requesting partition status on partition without lock. Topic: " << topic << ". Partition: " << partition, "", SessionId); + } else if (lockIt->second.Gen != generation) { + WARN_LOG("Requesting partition status on partition with wrong generation of lock. Topic: " << topic << ". Partition: " << partition << ". Generation: " << generation, "", SessionId); + } else { + Consumer->RequestPartitionStatus(topic, partition, lockIt->second.OriginalGen); + } +} + +void TRetryingConsumer::Cancel() { + Destroy(GetCancelReason()); +} + +NThreading::TFuture<void> TRetryingConsumer::Destroyed() noexcept { + return ConsumerDestroyedPromise.GetFuture(); +} + +void TRetryingConsumer::OnStartDeadline() { + if (!StartPromise.HasValue()) { + TError error; + error.SetDescription("Start timeout."); + error.SetCode(NErrorCode::CREATE_TIMEOUT); + Destroy(error); + } +} + +void TRetryingConsumer::OnConsumerDead(const TError& error) { + WARN_LOG("Subconsumer is dead: " << error, "", SessionId); + ScheduleReconnect(); +} + +void TRetryingConsumer::ScheduleReconnect() { + if (Stopping || Reconnecting) { + return; + } + Reconnecting = true; + if (!Locks.empty()) { + // need to notify client that all locks are expired + for (auto&& [key, lock] : Locks) { + if (lock.Locked) { + TReadResponse response; + auto* release = response.MutableRelease(); + release->set_generation(lock.Gen); + release->set_topic(key.first); + release->set_partition(key.second); + release->set_can_commit(false); + FastResponse(TConsumerMessage(std::move(response))); + } + } + } + if (!CommittingCookies.empty()) { + // need to notify client that all cookies are commited, because these cookies aren't valid anymore + TReadResponse response; + auto* cookies = response.MutableCommit()->MutableCookie(); + cookies->Reserve(CommittingCookies.size()); + for (const auto& cookiePair : CommittingCookies) { + cookies->Add(cookiePair.second); + } + FastResponse(TConsumerMessage(std::move(response))); + } + + Cookies.clear(); + CommittingCookies.clear(); + Locks.clear(); + Consumer = nullptr; + + if (ReconnectionAttemptsDone >= Settings.MaxAttempts) { + Destroy(TStringBuilder() << "Failed " << ReconnectionAttemptsDone << " reconnection attempts"); + return; + } + + ++ReconnectionAttemptsDone; + + TDuration delay = Min(Settings.MaxReconnectionDelay, ReconnectionAttemptsDone * Settings.ReconnectionDelay); + std::weak_ptr<TRetryingConsumer> weakRef = shared_from_this(); + PQLib->GetScheduler().Schedule(delay, this, [weakRef]() { + auto self = weakRef.lock(); + if (nullptr != self) { + self->DoReconnect(self->Settings.StartSessionTimeout.ToDeadLine()); + self->Reconnecting = false; + } + }); +} + +void TRetryingConsumer::DoReconnect(TInstant deadline) { + if (Stopping) { + return; + } + DEBUG_LOG("Create subconsumer", "", SessionId); + Consumer = PQLib->CreateRawConsumer(Settings, DestroyEventRef, Logger); + std::weak_ptr<TRetryingConsumer> weak(shared_from_this()); + StartFuture = Consumer->Start(deadline); + PQLib->Subscribe(StartFuture, this, [weak](const NThreading::TFuture<TConsumerCreateResponse>& f) { + auto self = weak.lock(); + if (nullptr != self) { + self->StartProcessing(f); + } + }); + + PQLib->Subscribe(Consumer->IsDead(), this, [weak](const NThreading::TFuture<TError>& error) { + auto self = weak.lock(); + if (nullptr != self) { + self->OnConsumerDead(error.GetValue()); + } + }); +} + +void TRetryingConsumer::StartProcessing(const NThreading::TFuture<TConsumerCreateResponse>& f) { + Y_VERIFY(f.HasValue()); + if (f.GetValue().Response.HasError()) { + WARN_LOG("Cannot create subconsumer: " << f.GetValue().Response.GetError(), "", SessionId); + if (IsRetryable(f.GetValue().Response.GetError())) { + ScheduleReconnect(); + } else { + Destroy(f.GetValue().Response.GetError()); + } + } else { + if (!SessionId) { + SessionId = f.GetValue().Response.GetInit().GetSessionId(); + } + ReconnectionAttemptsDone = 0; + // need to schedule again all pending requests which are running before connection lost + for (size_t cnt = PendingRequests.size(); 0 != cnt; --cnt) { + DoRequest(); + } + StartPromise.TrySetValue(f.GetValue()); + } +} + +void TRetryingConsumer::SubscribeDestroyed() { + NThreading::TPromise<void> promise = ConsumerDestroyedPromise; + auto handler = [promise](const auto&) mutable { + promise.SetValue(); + }; + if (Consumer) { + WaitExceptionOrAll(DestroyedPromise.GetFuture(), Consumer->Destroyed()).Subscribe(handler); + } else { + DestroyedPromise.GetFuture().Subscribe(handler); + } + + DestroyPQLibRef(); +} + +void TRetryingConsumer::Destroy(const TError& error) { + if (Stopping) { + return; + } + + Stopping = true; + + SubscribeDestroyed(); + + TConsumerMessage message = [&]() { + TReadResponse response; + *response.MutableError() = error; + return TConsumerMessage(std::move(response)); + }(); + + for (auto& r : PendingRequests) { + r.TrySetValue(message); + } + IsDeadPromise.TrySetValue(error); + if (StartPromise.Initialized()) { + StartPromise.TrySetValue(TConsumerCreateResponse(std::move(message.Response))); + } + PendingRequests.clear(); + ReadyResponses.clear(); + Cookies.clear(); + CommittingCookies.clear(); + Locks.clear(); + Consumer = nullptr; + Reconnecting = false; +} + +void TRetryingConsumer::Destroy(const TString& description, NErrorCode::EErrorCode code) { + TError error; + error.SetDescription(description); + error.SetCode(code); + Destroy(error); +} + +void TRetryingConsumer::DoRequest() { + Y_VERIFY(Consumer); + std::weak_ptr<TRetryingConsumer> weak(shared_from_this()); + PQLib->Subscribe(Consumer->GetNextMessage(), this, [weak](NThreading::TFuture<TConsumerMessage>& f) mutable { + auto self = weak.lock(); + if (nullptr != self) { + self->ProcessResponse(f.ExtractValueSync()); + } + }); +} + +void TRetryingConsumer::ProcessResponse(TConsumerMessage&& message) { + switch (message.Type) { + case NPersQueue::EMessageType::EMT_DATA: + { + auto* data = message.Response.mutable_data(); + TCookieInfo* cookie = nullptr; + if (!Settings.CommitsDisabled) { + Cookies.push_back(TCookieInfo{{}, data->cookie(), CookieCounter++}); + cookie = &Cookies.back(); + data->set_cookie(cookie->UserCookie); + DEBUG_LOG("Got data from subconsumer. Cookie: " << cookie->OriginalCookie << ". User cookie: " << cookie->UserCookie, "", SessionId); + } + + for (auto& b : data->message_batch()) { + auto lockIt = Locks.find(std::make_pair(b.topic(), b.partition())); + if (lockIt != Locks.end()) { + Y_VERIFY(lockIt->second.Locked); + if (!Settings.CommitsDisabled) { + Y_ASSERT(nullptr != cookie); + cookie->Locks.emplace(&lockIt->second); + lockIt->second.Cookies.emplace(cookie->UserCookie); + } + + // Validate that all offsets are >= LockInfo.ReadOffset as expected. + if (lockIt->second.ReadOffset != 0) { + for (const auto& message : b.message()) { + if (message.offset() < lockIt->second.ReadOffset) { + Destroy( + TStringBuilder() + << "Fatal error: expected offsets for topic " << b.topic() + << " partition " << b.partition() << " >= " << lockIt->second.ReadOffset + << ", but got offset " << message.offset() + ); + return; + } + } + } + } + } + // TODO (bulatman) what about batched_data? + } + break; + case NPersQueue::EMessageType::EMT_LOCK: + { + auto* lock = message.Response.mutable_lock(); + auto& lockInfo = Locks[std::make_pair(lock->topic(), lock->partition())]; + Y_VERIFY(!lockInfo.Locked); + lockInfo.Locked = true; + lockInfo.Gen = GenCounter++; + lockInfo.OriginalGen = lock->generation(); + + lock->set_generation(lockInfo.Gen); + + DEBUG_LOG("Got lock from subconsumer on (" << lock->topic() << ", " << lock->partition() << "). Generation: " << lockInfo.OriginalGen << ". User generation: " << lockInfo.Gen, "", SessionId); + + std::weak_ptr<TRetryingConsumer> weak = shared_from_this(); + PQLib->Subscribe(message.ReadyToRead.GetFuture(), this, + [weak, topic = lock->topic(), partition = lock->partition(), generation = lock->generation()](const NThreading::TFuture<NPersQueue::TLockInfo>& lockInfo) { + auto self = weak.lock(); + if (nullptr != self) { + self->UpdateReadyToRead(lockInfo.GetValue(), topic, partition, generation); + } + }); + } + break; + case NPersQueue::EMessageType::EMT_RELEASE: + { + auto* release = message.Response.mutable_release(); + const bool softRelease = release->can_commit(); + auto lockIt = Locks.find(std::make_pair(release->topic(), release->partition())); + const bool lockInfoFound = lockIt != Locks.end(); + DEBUG_LOG("Got release from subconsumer on (" << release->topic() << ", " << release->partition() << "). Can commit: " << softRelease << ". Has lock info: " << lockInfoFound << ". Generation: " << release->generation() << ". User generation: " << (lockInfoFound ? lockIt->second.Gen : 0), "", SessionId); + if (lockInfoFound) { // It is normal situation when client receives Release(canCommit=true) and then Release(canCommit=false). + auto& lockInfo = lockIt->second; + Y_VERIFY(lockInfo.OriginalGen == release->generation(), "lock generation mismatch"); + release->set_generation(lockInfo.Gen); + + if (softRelease) { + lockInfo.Locked = false; + } else { + for (auto cookie : lockInfo.Cookies) { + auto& cookieInfo = Cookies[cookie + Cookies.size() - CookieCounter]; + Y_VERIFY(cookieInfo.UserCookie == cookie); + // Lock is not valid anymore. + cookieInfo.Locks.erase(&lockInfo); + } + Locks.erase(lockIt); + } + } else { + return; + } + } + break; + case NPersQueue::EMessageType::EMT_STATUS: + { + auto* status = message.Response.mutable_partition_status(); + auto lockIt = Locks.find(std::make_pair(status->topic(), status->partition())); + Y_VERIFY(lockIt != Locks.end() && lockIt->second.Locked && lockIt->second.OriginalGen == status->generation()); + status->set_generation(lockIt->second.Gen); + } + break; + case NPersQueue::EMessageType::EMT_COMMIT: + { + auto* cookies = message.Response.mutable_commit()->mutable_cookie(); + // convert cookies + for (int i = 0; i < cookies->size(); ++i) { + auto it = CommittingCookies.find(cookies->Get(i)); + Y_VERIFY(it != CommittingCookies.end(), "unknown commited cookie!"); + cookies->Set(i, it->second); + CommittingCookies.erase(it); + } + } + break; + + case NPersQueue::EMessageType::EMT_ERROR: + // check error, if retryable, need to recreate consumer and read message again + const bool retryable = IsRetryable(message.Response.error()); + DEBUG_LOG("Got error from subconsumer: " << message.Response.error() << ". Retryable: " << retryable, "", SessionId); + if (retryable) { + ScheduleReconnect(); + return; + } + break; + }; + + if (!PendingRequests.empty()) { + PendingRequests.front().SetValue(std::move(message)); + PendingRequests.pop_front(); + } else { + ReadyResponses.push_back(std::move(message)); + } +} + +void TRetryingConsumer::FastResponse(TConsumerMessage&& message) { + if (!PendingRequests.empty()) { + PendingRequests.front().SetValue(std::move(message)); + PendingRequests.pop_front(); + } else { + ReadyResponses.push_back(std::move(message)); + } +} + +void TRetryingConsumer::FastCommit(const TVector<ui64>& cookies) { + TReadResponse response; + auto* cookiesMessage = response.MutableCommit()->MutableCookie(); + cookiesMessage->Reserve(cookies.size()); + for (auto cookie : cookies) { + cookiesMessage->Add(cookie); + } + FastResponse(TConsumerMessage(std::move(response))); +} + +void TRetryingConsumer::UpdateReadyToRead(const NPersQueue::TLockInfo& readyToRead, const TString& topic, ui32 partition, ui64 generation) { + const auto lockIt = Locks.find(std::make_pair(topic, partition)); + if (lockIt != Locks.end() && lockIt->second.Locked && lockIt->second.Gen == generation) { + lockIt->second.ReadOffset = readyToRead.ReadOffset; + } +} + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_consumer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_consumer.h new file mode 100644 index 0000000000..448d35e4c8 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_consumer.h @@ -0,0 +1,107 @@ +#pragma once + +#include "consumer.h" +#include "scheduler.h" +#include "internals.h" +#include "persqueue_p.h" +#include "iconsumer_p.h" + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include <kikimr/yndx/api/grpc/persqueue.grpc.pb.h> + +#include <library/cpp/threading/future/future.h> + +#include <util/generic/hash.h> +#include <util/generic/hash_set.h> +#include <util/generic/vector.h> + +#include <deque> + +namespace NPersQueue { + +// @brief consumer implementation which transparently retries all connectivity errors +class TRetryingConsumer: public IConsumerImpl, public std::enable_shared_from_this<TRetryingConsumer> { + // @brief locked partitions info + struct TLockInfo { + THashSet<ui64> Cookies; // related cookies + ui64 Gen = 0; + ui64 OriginalGen = 0; // original generation + ui64 ReadOffset = 0; // read offset specified by client + bool Locked = false; + }; + + struct TCookieInfo { + THashSet<TLockInfo*> Locks; // related locks + ui64 OriginalCookie = 0; // zero means invalid cookie + ui64 UserCookie = 0; + }; + +public: + TRetryingConsumer(const TConsumerSettings& settings, std::shared_ptr<void> destroyEventRef, + TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger); + + ~TRetryingConsumer() noexcept override; + + NThreading::TFuture<TConsumerCreateResponse> Start(TInstant deadline = TInstant::Max()) noexcept override; + + using IConsumerImpl::GetNextMessage; + void GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept override; + + void Commit(const TVector<ui64>& cookies) noexcept override; + + void RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept override; + + NThreading::TFuture<TError> IsDead() noexcept override; + + void Cancel() override; + + NThreading::TFuture<void> Destroyed() noexcept override; + +private: + void OnStartDeadline(); + void OnConsumerDead(const TError& error); + void ScheduleReconnect(); + void DoReconnect(TInstant deadline); + void StartProcessing(const NThreading::TFuture<TConsumerCreateResponse>& f); + void SubscribeDestroyed(); + void Destroy(const TError& error); + void Destroy(const TString& description, NErrorCode::EErrorCode code = NErrorCode::ERROR); + void DoRequest(); + void ProcessResponse(TConsumerMessage&& message); + void FastResponse(TConsumerMessage&& message); + void FastCommit(const TVector<ui64>& cookies); + void UpdateReadyToRead(const NPersQueue::TLockInfo& readyToRead, const TString& topic, ui32 partition, ui64 generation); + +private: + TConsumerSettings Settings; + TIntrusivePtr<ILogger> Logger; + std::shared_ptr<IConsumerImpl> Consumer; + TString SessionId; + + NThreading::TFuture<TConsumerCreateResponse> StartFuture; + NThreading::TPromise<TConsumerCreateResponse> StartPromise; + NThreading::TPromise<TError> IsDeadPromise; + NThreading::TPromise<void> ConsumerDestroyedPromise; + + // requests which are waiting for response + std::deque<NThreading::TPromise<TConsumerMessage>> PendingRequests; + // ready messages for returning to clients immediately + std::deque<TConsumerMessage> ReadyResponses; + // active cookies + std::deque<TCookieInfo> Cookies; + THashMap<ui64, ui64> CommittingCookies; + // active data per topic, partition + THashMap<std::pair<TString, ui64>, TLockInfo> Locks; // topic, partition -> LockInfo + + // number of unsuccessful retries after last error + ui64 GenCounter; + ui64 CookieCounter; + unsigned ReconnectionAttemptsDone; + // destroying process started + bool Stopping; + // reconnecting in process + bool Reconnecting; +}; +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_consumer_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_consumer_ut.cpp new file mode 100644 index 0000000000..8065ed717a --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_consumer_ut.cpp @@ -0,0 +1,604 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/sdk_test_setup.h> + +#include <library/cpp/testing/unittest/registar.h> + +#include <util/generic/guid.h> + +using namespace NThreading; +using namespace NKikimr; +using namespace NKikimr::NPersQueueTests; + +namespace NPersQueue { + Y_UNIT_TEST_SUITE(TRetryingConsumerTest) { + static TConsumerSettings FakeSettings() { + TConsumerSettings settings; + settings.ReconnectOnFailure = true; + settings.Server = TServerSetting{"localhost"}; + settings.Topics.push_back("topic"); + settings.ClientId = "client"; + return settings; + } + + Y_UNIT_TEST(NotStartedConsumerCanBeDestructed) { + // Test that consumer doesn't hang on till shutdown + TPQLib lib; + auto settings = FakeSettings(); + lib.CreateConsumer(settings); + } + + Y_UNIT_TEST(StartDeadlineExpires) { + TPQLibSettings libSettings; + libSettings.ChannelCreationTimeout = TDuration::MilliSeconds(1); + libSettings.DefaultLogger = new TCerrLogger(TLOG_DEBUG); + TPQLib lib(libSettings); + THolder<IConsumer> consumer; + { + auto settings = FakeSettings(); + settings.ReconnectionDelay = TDuration::MilliSeconds(1); + settings.StartSessionTimeout = TDuration::MilliSeconds(10); + if (GrpcV1EnabledByDefault()) { + settings.MaxAttempts = 3; + } + consumer = lib.CreateConsumer(settings); + } + const TInstant beforeStart = TInstant::Now(); + auto future = consumer->Start(TDuration::MilliSeconds(100)); + if (!GrpcV1EnabledByDefault()) { + UNIT_ASSERT(future.GetValueSync().Response.HasError()); + UNIT_ASSERT_EQUAL_C(future.GetValueSync().Response.GetError().GetCode(), NErrorCode::CREATE_TIMEOUT, + "Error: " << future.GetValueSync().Response.GetError()); + const TInstant now = TInstant::Now(); + UNIT_ASSERT_C(now - beforeStart >= TDuration::MilliSeconds(100), now); + } + auto isDead = consumer->IsDead(); + isDead.Wait(); + + DestroyAndWait(consumer); + } + + Y_UNIT_TEST(StartMaxAttemptsExpire) { + TPQLibSettings libSettings; + libSettings.ChannelCreationTimeout = TDuration::MilliSeconds(1); + libSettings.DefaultLogger = new TCerrLogger(TLOG_DEBUG); + TPQLib lib(libSettings); + THolder<IConsumer> consumer; + { + auto settings = FakeSettings(); + settings.ReconnectionDelay = TDuration::MilliSeconds(10); + settings.StartSessionTimeout = TDuration::MilliSeconds(10); + settings.MaxAttempts = 3; + consumer = lib.CreateConsumer(settings); + } + const TInstant beforeStart = TInstant::Now(); + auto future = consumer->Start(); + if (!GrpcV1EnabledByDefault()) { + UNIT_ASSERT(future.GetValueSync().Response.HasError()); + } + auto deadFuture = consumer->IsDead(); + deadFuture.Wait(); + const TInstant now = TInstant::Now(); + UNIT_ASSERT_C(now - beforeStart >= TDuration::MilliSeconds(30), now); + + DestroyAndWait(consumer); + } + + static void AssertWriteValid(const NThreading::TFuture<TProducerCommitResponse>& respFuture) { + const TProducerCommitResponse& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TWriteResponse::kAck, "Msg: " << resp.Response); + } + + static void AssertReadValid(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kData, "Msg: " << resp.Response); + } + + static void AssertCommited(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kCommit, "Msg: " << resp.Response); + } + + static void AssertReadFailed(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kError, "Msg: " << resp.Response); + } + + static TProducerSettings MakeProducerSettings(const TTestServer& testServer) { + TProducerSettings producerSettings; + producerSettings.ReconnectOnFailure = false; + producerSettings.Topic = "topic1"; + producerSettings.SourceId = "123"; + producerSettings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + producerSettings.Codec = ECodec::LZOP; + return producerSettings; + } + + static TConsumerSettings MakeConsumerSettings(const TTestServer& testServer) { + TConsumerSettings settings; + settings.ReconnectOnFailure = true; + settings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + settings.Topics.emplace_back("topic1"); + settings.ClientId = "user"; + return settings; + } + + Y_UNIT_TEST(ReconnectsToServer) { + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(false); + + testServer.AnnoyingClient->FullInit(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + const size_t partitions = 10; + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", partitions); + + testServer.WaitInit("topic1"); + + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLib PQLib; + + auto producer = PQLib.CreateProducer(MakeProducerSettings(testServer), logger, false); + UNIT_ASSERT(!producer->IsDead().HasValue()); + producer->Start().Wait(); + + auto consumer = PQLib.CreateConsumer(MakeConsumerSettings(testServer), logger, false); + TFuture<TError> isDead = consumer->IsDead(); + UNIT_ASSERT(!isDead.HasValue()); + consumer->Start().Wait(); + + auto read1 = consumer->GetNextMessage(); + + // write first value + auto write1 = producer->Write(1, TString("blob1")); + AssertWriteValid(write1); + AssertReadValid(read1); + + // commit + consumer->Commit({1}); + auto commitAck = consumer->GetNextMessage(); + Cerr << "wait for commit1\n"; + + AssertCommited(commitAck); + + Cerr << "After shutdown\n"; + testServer.ShutdownGRpc(); + + auto read2 = consumer->GetNextMessage(); + UNIT_ASSERT(!isDead.HasValue()); + testServer.EnableGRpc(); + + testServer.WaitInit("topic1"); + + Cerr << "Wait producer1 death\n"; + + producer->IsDead().Wait(); + + producer = PQLib.CreateProducer(MakeProducerSettings(testServer), logger, false); + Cerr << "Wait producer2 start\n"; + producer->Start().Wait(); + UNIT_ASSERT(!producer->IsDead().HasValue()); + auto write2 = producer->Write(2, TString("blob2")); + Cerr << "Wait for second write\n"; + AssertWriteValid(write2); + Cerr << "wait for read second blob\n"; + AssertReadValid(read2); + UNIT_ASSERT_VALUES_EQUAL(2u, read2.GetValueSync().Response.GetData().GetCookie()); + UNIT_ASSERT(!isDead.HasValue()); + consumer->Commit({2}); + Cerr << "wait for commit2\n"; + auto commitAck2 = consumer->GetNextMessage(); + AssertCommited(commitAck2); + + DestroyAndWait(producer); + DestroyAndWait(consumer); + } + + static void DiesOnTooManyReconnectionAttempts(bool callRead) { + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(false); + + testServer.AnnoyingClient->FullInit(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", 2); + + testServer.WaitInit("topic1"); + + TConsumerSettings settings = MakeConsumerSettings(testServer); + settings.MaxAttempts = 3; + settings.ReconnectionDelay = TDuration::MilliSeconds(100); + auto consumer = testServer.PQLib->CreateConsumer(settings, testServer.PQLibSettings.DefaultLogger, false); + consumer->Start().Wait(); + TFuture<TError> isDead = consumer->IsDead(); + UNIT_ASSERT(!isDead.HasValue()); + + // shutdown server + const TInstant beforeShutdown = TInstant::Now(); + testServer.ShutdownServer(); + + NThreading::TFuture<TConsumerMessage> read; + if (callRead) { + read = consumer->GetNextMessage(); + } + + isDead.Wait(); + const TInstant afterDead = TInstant::Now(); + // 3 attempts: 100ms, 200ms and 300ms + UNIT_ASSERT_C(afterDead - beforeShutdown >= TDuration::MilliSeconds(GrpcV1EnabledByDefault() ? 300 : 600), "real difference: " << (afterDead - beforeShutdown)); + + if (callRead) { + AssertReadFailed(read); + } + } + + Y_UNIT_TEST(DiesOnTooManyReconnectionAttemptsWithoutRead) { + // Check that we reconnect even without explicit write errors + DiesOnTooManyReconnectionAttempts(false); + } + + Y_UNIT_TEST(DiesOnTooManyReconnectionAttemptsWithRead) { + DiesOnTooManyReconnectionAttempts(true); + } + + Y_UNIT_TEST(ReadBeforeConnectToServer) { + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(false); + + testServer.AnnoyingClient->FullInit(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + const size_t partitions = 10; + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", partitions); + + testServer.WaitInit("topic1"); + + auto consumer = testServer.PQLib->CreateConsumer( + MakeConsumerSettings(testServer), testServer.PQLibSettings.DefaultLogger, false + ); + auto consumerStarted = consumer->Start(); + auto read1 = consumer->GetNextMessage(); + consumerStarted.Wait(); + TVector<TFuture<TConsumerMessage>> messages; + for (int i = 2; i <= 5; ++i) { + messages.push_back(consumer->GetNextMessage()); + } + + auto producer = testServer.PQLib->CreateProducer( + MakeProducerSettings(testServer), testServer.PQLibSettings.DefaultLogger, false + ); + producer->Start().Wait(); + + while (!messages.back().HasValue()) { + auto write = producer->Write(TString("data")); + AssertWriteValid(write); + } + + AssertReadValid(read1); + for (auto& msg: messages) { + AssertReadValid(msg); + } + + DestroyAndWait(producer); + DestroyAndWait(consumer); + } + + Y_UNIT_TEST(CancelsStartAfterPQLibDeath) { + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLibSettings pqLibSettings; + pqLibSettings.DefaultLogger = logger; + THolder<TPQLib> PQLib = MakeHolder<TPQLib>(pqLibSettings); + + auto consumer = PQLib->CreateConsumer(FakeSettings(), logger, false); + auto start = consumer->Start(); + if (!GrpcV1EnabledByDefault()) { + UNIT_ASSERT(!start.HasValue()); + } + + PQLib = nullptr; + + if (!GrpcV1EnabledByDefault()) { + UNIT_ASSERT(start.HasValue()); + UNIT_ASSERT(start.GetValue().Response.HasError()); + } + + auto dead = consumer->IsDead(); + if (GrpcV1EnabledByDefault()) { + dead.Wait(); + } else { + UNIT_ASSERT(dead.HasValue()); + } + + auto read = consumer->GetNextMessage(); + UNIT_ASSERT(read.HasValue()); + UNIT_ASSERT(read.GetValueSync().Response.HasError()); + } + + Y_UNIT_TEST(LockSession) { + SDKTestSetup setup("LockSession", false); + setup.GetGrpcServerOptions().SetGRpcShutdownDeadline(TDuration::MilliSeconds(100)); + setup.Start(); + auto consumerSettings = setup.GetConsumerSettings(); + consumerSettings.UseLockSession = true; + consumerSettings.ReconnectOnFailure = true; + auto consumer = setup.StartConsumer(consumerSettings); + + setup.WriteToTopic({"msg1", "msg2"}); + + size_t releases = 0; + for (size_t times = 0; times < 2; ++times) { + while (true) { + auto msgFuture = consumer->GetNextMessage(); + auto& msg = msgFuture.GetValueSync(); + setup.GetLog() << TLOG_INFO << "Got response from consumer: " << msg.Response; + if (times) { + UNIT_ASSERT_C(msg.Type == EMT_LOCK || msg.Type == EMT_DATA || msg.Type == EMT_RELEASE, msg.Response); + } else { + UNIT_ASSERT_C(msg.Type == EMT_LOCK || msg.Type == EMT_DATA, msg.Response); + } + if (msg.Type == EMT_LOCK) { + TLockInfo lockInfo; + lockInfo.ReadOffset = 1; + msg.ReadyToRead.SetValue(lockInfo); + } else if (msg.Type == EMT_DATA) { + UNIT_ASSERT_VALUES_EQUAL_C(msg.Response.GetData().MessageBatchSize(), 1, msg.Response); + UNIT_ASSERT_VALUES_EQUAL_C(msg.Response.GetData().GetMessageBatch(0).MessageSize(), 1, msg.Response); + UNIT_ASSERT_VALUES_EQUAL_C(msg.Response.GetData().GetMessageBatch(0).GetMessage(0).GetData(), "msg2", msg.Response); + break; + } else if (msg.Type == EMT_RELEASE) { + ++releases; + } + } + if (!times) { // force reconnect + setup.ShutdownGRpc(); + setup.EnableGRpc(); + } + } + UNIT_ASSERT_VALUES_EQUAL(releases, 1); + } + + Y_UNIT_TEST(ReadAhead) { + SDKTestSetup setup("ReadAhead"); + + auto consumerSettings = setup.GetConsumerSettings(); + consumerSettings.ReconnectOnFailure = true; + auto consumer = setup.StartConsumer(consumerSettings); + + setup.WriteToTopic({"msg"}); + + // Wait until read ahead occurs. + Sleep(TDuration::MilliSeconds(100)); + + setup.ReadFromTopic({{"msg"}}, true, consumer.Get()); + } + + Y_UNIT_TEST(ReadBeforeWrite) { + SDKTestSetup setup("ReadBeforeWrite"); + + auto consumerSettings = setup.GetConsumerSettings(); + consumerSettings.ReconnectOnFailure = true; + auto consumer = setup.StartConsumer(consumerSettings); + + auto msgFuture = consumer->GetNextMessage(); + + Sleep(TDuration::MilliSeconds(100)); + + setup.WriteToTopic({"msg"}); + + const auto& msg = msgFuture.GetValueSync(); + UNIT_ASSERT_EQUAL(msg.Type, EMT_DATA); + UNIT_ASSERT_STRINGS_EQUAL(msg.Response.GetData().GetMessageBatch(0).GetMessage(0).GetData(), "msg"); + } + + Y_UNIT_TEST(CommitDisabled) { + SDKTestSetup setup("CommitDisabled"); + + auto consumerSettings = setup.GetConsumerSettings(); + consumerSettings.ReconnectOnFailure = true; + consumerSettings.CommitsDisabled = true; + auto consumer = setup.StartConsumer(consumerSettings); + setup.WriteToTopic({"msg", "msg2"}); + setup.ReadFromTopic({{"msg", "msg2"}}, false, consumer.Get()); + // check that commit fails with error when commits disabled + setup.WriteToTopic({"msg3"}); + auto msgFuture = consumer->GetNextMessage(); + const auto& msg = msgFuture.GetValueSync(); + UNIT_ASSERT_C(!msg.Response.HasError(), msg.Response); + UNIT_ASSERT_EQUAL(msg.Type, EMT_DATA); + UNIT_ASSERT_STRINGS_EQUAL(msg.Response.GetData().GetMessageBatch(0).GetMessage(0).GetData(), "msg3"); + consumer->Commit({msg.Response.GetData().GetCookie()}); + auto isDead = consumer->IsDead(); + auto& err = isDead.GetValueSync(); + UNIT_ASSERT_STRINGS_EQUAL("Commits are disabled", err.GetDescription()); + } + + void LockReleaseTest(bool multiclusterConsumer, bool killPqrb) { + SDKTestSetup setup("LockRelease"); + const TString topicName = "lock_release"; + const size_t partitionsCount = 10; + setup.CreateTopic(topicName, setup.GetLocalCluster(), partitionsCount); + + auto consumerSettings = setup.GetConsumerSettings(); + consumerSettings.Topics[0] = topicName; + consumerSettings.ReconnectOnFailure = true; + consumerSettings.UseLockSession = true; + if (multiclusterConsumer) { // The same test for multicluster wrapper. + consumerSettings.ReadFromAllClusterSources = true; + consumerSettings.ReadMirroredPartitions = false; + } + auto consumer = setup.StartConsumer(consumerSettings); + + std::vector<int> partitionsLocked(partitionsCount * 2); // For both consumers. + std::vector<ui64> partitionsGens(partitionsCount * 2); + for (size_t i = 0; i < partitionsCount; ++i) { + auto lock = consumer->GetNextMessage().GetValueSync(); + UNIT_ASSERT_EQUAL_C(lock.Type, EMT_LOCK, lock.Response); + UNIT_ASSERT_C(!partitionsLocked[lock.Response.GetLock().GetPartition()], lock.Response); + partitionsLocked[lock.Response.GetLock().GetPartition()] = 1; + partitionsGens[lock.Response.GetLock().GetPartition()] = lock.Response.GetLock().GetGeneration(); + lock.ReadyToRead.SetValue({}); + } + + const size_t messagesCount = 100; + auto generateMessages = [&]() { + for (size_t i = 0; i < messagesCount; ++i) { // Write to several random partitions. + const TString uniqStr = CreateGuidAsString(); + auto settings = setup.GetProducerSettings(); + settings.Topic = topicName; + settings.SourceId = uniqStr; + auto producer = setup.StartProducer(settings); + setup.WriteToTopic({uniqStr}, producer.Get()); + } + }; + + generateMessages(); + + THolder<IConsumer> otherConsumer = nullptr; + bool rebalanced = false; + + NThreading::TFuture<TConsumerMessage> msg1 = consumer->GetNextMessage(); + NThreading::TFuture<TConsumerMessage> msg2; + size_t hardReleasesCount = 0; + + THashSet<TString> msgsReceived; + THashSet<ui64> cookiesToCommit; + auto locksCount = [&]() -> int { + return std::accumulate(partitionsLocked.begin(), partitionsLocked.end(), 0); + }; + auto allEventsReceived = [&]() -> bool { + return msgsReceived.size() >= messagesCount * 2 + && locksCount() == partitionsCount + && (!killPqrb || hardReleasesCount > 0) + && cookiesToCommit.empty(); + }; + + int iterationNum = 0; + while (!allEventsReceived()) { + ++iterationNum; + Cerr << "\nIterationNum: " << iterationNum << Endl; + Cerr << "msgsReceived.size(): " << msgsReceived.size() << Endl; + Cerr << "partitionsCount: " << partitionsCount << Endl; + Cerr << "locksGot: " << locksCount() << Endl; + Cerr << "hardReleasesCount: " << hardReleasesCount << Endl; + Cerr << "cookiesToCommit.size(): " << cookiesToCommit.size() << Endl; + auto waiter = msg2.Initialized() ? NThreading::WaitAny(msg1.IgnoreResult(), msg2.IgnoreResult()) : msg1.IgnoreResult(); + waiter.Wait(); + auto processValue = [&](auto& msg, auto& consumer) { + const bool isOtherConsumer = consumer == otherConsumer; + const ui64 consumerMix = isOtherConsumer ? 0x8000000000000000 : 0; + + switch (msg.GetValue().Type) { + case EMT_LOCK: + { + UNIT_ASSERT_LT(msg.GetValue().Response.GetLock().GetPartition(), partitionsCount); + const size_t idx = msg.GetValue().Response.GetLock().GetPartition() + (isOtherConsumer ? partitionsCount : 0); + if (msg.GetValue().Response.GetLock().GetGeneration() >= partitionsGens[idx]) { + UNIT_ASSERT_C(!partitionsLocked[idx], msg.GetValue().Response << ". Other: " << isOtherConsumer); + partitionsLocked[idx] = 1; + msg.GetValue().ReadyToRead.SetValue({}); + partitionsGens[idx] = msg.GetValue().Response.GetLock().GetGeneration(); + } + break; + } + case EMT_RELEASE: + { + UNIT_ASSERT_LT(msg.GetValue().Response.GetRelease().GetPartition(), partitionsCount); + const size_t idx = msg.GetValue().Response.GetRelease().GetPartition() + (isOtherConsumer ? partitionsCount : 0); + partitionsLocked[idx] = 0; + if (!msg.GetValue().Response.GetRelease().GetCanCommit()) { + ++hardReleasesCount; + } + if (!killPqrb) { + UNIT_ASSERT(msg.GetValue().Response.GetRelease().GetCanCommit()); // No restarts => soft release. + } + break; + } + case EMT_DATA: + { + const auto& resp = msg.GetValue().Response.GetData(); + UNIT_ASSERT(cookiesToCommit.insert(resp.GetCookie() | consumerMix).second); + for (const auto& batch : resp.GetMessageBatch()) { + for (const auto& message : batch.GetMessage()) { + msgsReceived.insert(message.GetData()); + } + } + consumer->Commit({resp.GetCookie()}); + break; + } + case EMT_ERROR: + case EMT_STATUS: + case EMT_COMMIT: + for (ui64 cookie : msg.GetValue().Response.GetCommit().GetCookie()) { + UNIT_ASSERT(cookiesToCommit.find(cookie | consumerMix) != cookiesToCommit.end()); + cookiesToCommit.erase(cookie | consumerMix); + } + break; + } + msg = consumer->GetNextMessage(); + }; + + if (msg1.HasValue()) { + processValue(msg1, consumer); + } + if (msg2.Initialized() && msg2.HasValue()) { + processValue(msg2, otherConsumer); + } + + if (!otherConsumer && msgsReceived.size() >= messagesCount / 4) { + otherConsumer = setup.StartConsumer(consumerSettings); + msg2 = otherConsumer->GetNextMessage(); + } + + if (!rebalanced && msgsReceived.size() >= messagesCount / 2) { + rebalanced = true; + + if (killPqrb) { + Sleep(TDuration::MilliSeconds(100)); + setup.KillPqrb(topicName, setup.GetLocalCluster()); + } + + generateMessages(); + } + } + + UNIT_ASSERT_VALUES_EQUAL(locksCount(), partitionsCount); + if (killPqrb) { + UNIT_ASSERT(hardReleasesCount > 0); + } else { + UNIT_ASSERT_VALUES_EQUAL(hardReleasesCount, 0); + } + } + + Y_UNIT_TEST(LockRelease) { + LockReleaseTest(false, false); + } + + Y_UNIT_TEST(LockReleaseHardRelease) { + LockReleaseTest(false, true); + } + + Y_UNIT_TEST(LockReleaseMulticluster) { + LockReleaseTest(true, false); + } + + Y_UNIT_TEST(LockReleaseMulticlusterHardRelease) { + LockReleaseTest(true, true); + } + + Y_UNIT_TEST(NonRetryableErrorOnStart) { + SDKTestSetup setup("NonRetryableErrorOnStart"); + + auto consumerSettings = setup.GetConsumerSettings(); + //consumerSettings. + consumerSettings.ReconnectOnFailure = true; + consumerSettings.Topics = { "unknown/topic" }; + THolder<IConsumer> consumer = setup.GetPQLib()->CreateConsumer(consumerSettings); + auto startFuture = consumer->Start(); + if (!GrpcV1EnabledByDefault()) { + Cerr << "===Got response: " << startFuture.GetValueSync().Response << Endl; + UNIT_ASSERT_C(startFuture.GetValueSync().Response.HasError(), startFuture.GetValueSync().Response); + UNIT_ASSERT_EQUAL(startFuture.GetValueSync().Response.GetError().GetCode(), NPersQueue::NErrorCode::UNKNOWN_TOPIC); + } + auto deadFuture = consumer->IsDead(); + UNIT_ASSERT_EQUAL_C(deadFuture.GetValueSync().GetCode(), NPersQueue::NErrorCode::UNKNOWN_TOPIC, deadFuture.GetValueSync()); + } + } +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_producer.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_producer.cpp new file mode 100644 index 0000000000..bb562ea0ae --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_producer.cpp @@ -0,0 +1,370 @@ +#include "retrying_producer.h" +#include "persqueue_p.h" + +#include <util/generic/strbuf.h> +#include <util/stream/zlib.h> +#include <util/string/cast.h> +#include <util/string/printf.h> +#include <util/string/vector.h> +#include <util/string/builder.h> + +namespace NPersQueue { + +TRetryingProducer::TRetryingProducer(const TProducerSettings& settings, std::shared_ptr<void> destroyEventRef, + TIntrusivePtr<TPQLibPrivate> pqLib, TIntrusivePtr<ILogger> logger) + : IProducerImpl(std::move(destroyEventRef), std::move(pqLib)) + , Settings(settings) + , Logger(std::move(logger)) + , IsDeadPromise(NThreading::NewPromise<TError>()) + , NeedRecreation(true) + , Stopping(false) + , ToProcess(0) + , LastReconnectionDelay(TDuration::Zero()) + , ReconnectionAttemptsDone(0) +{ + if (Settings.MaxAttempts == 0) { + ythrow yexception() << "MaxAttempts setting can't be zero."; + } + + if (Settings.ReconnectionDelay == TDuration::Zero()) { + ythrow yexception() << "ReconnectionDelay setting can't be zero."; + } + + if (Settings.StartSessionTimeout < PQLib->GetSettings().ChannelCreationTimeout) { + ythrow yexception() << "StartSessionTimeout can't be less than ChannelCreationTimeout."; + } +} + +TRetryingProducer::~TRetryingProducer() noexcept { + Destroy("Destructor called"); +} + +NThreading::TFuture<TProducerCreateResponse> TRetryingProducer::Start(TInstant deadline) noexcept { + StartPromise = NThreading::NewPromise<TProducerCreateResponse>(); + if (deadline != TInstant::Max()) { + std::weak_ptr<TRetryingProducer> self = shared_from_this(); + auto onDeadline = [self] { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->OnStartDeadline(); + } + }; + PQLib->GetScheduler().Schedule(deadline, this, onDeadline); + } + RecreateProducer(deadline); + return StartPromise.GetFuture(); +} + +void TRetryingProducer::RecreateProducer(TInstant deadline) noexcept { + Producer = nullptr; + if (Stopping) + return; + + Y_VERIFY(InFlightRequests.size() == Futures.size()); + DEBUG_LOG("Recreating subproducer. Futures size: " << Futures.size(), Settings.SourceId, ""); + if (Futures.empty()) { + DoRecreate(deadline); + } // otherwise it will be recreated in ProcessFutures() +} + +void TRetryingProducer::DoRecreate(TInstant deadline) noexcept { + Y_VERIFY(InFlightRequests.size() == Futures.size()); + Y_VERIFY(Futures.empty()); + if (Stopping) + return; + Y_VERIFY(NeedRecreation); + NeedRecreation = false; + if (ReconnectionAttemptsDone >= Settings.MaxAttempts) { + Destroy(TStringBuilder() << "Failed " << ReconnectionAttemptsDone << " reconnection attempts"); + return; + } + ++ReconnectionAttemptsDone; + DEBUG_LOG("Creating subproducer. Attempt: " << ReconnectionAttemptsDone, Settings.SourceId, ""); + Producer = PQLib->CreateRawProducer(Settings, DestroyEventRef, Logger); + StartFuture = Producer->Start(deadline); + + std::weak_ptr<TRetryingProducer> self(shared_from_this()); + PQLib->Subscribe(StartFuture, + this, + [self](const NThreading::TFuture<TProducerCreateResponse>& f) { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->ProcessStart(f); + } + }); +} + +TDuration TRetryingProducer::UpdateReconnectionDelay() { + if (LastReconnectionDelay == TDuration::Zero()) { + LastReconnectionDelay = Settings.ReconnectionDelay; + } else { + LastReconnectionDelay *= 2; + } + LastReconnectionDelay = Min(LastReconnectionDelay, Settings.MaxReconnectionDelay); + return LastReconnectionDelay; +} + +void TRetryingProducer::ScheduleRecreation() { + if (Stopping) { + return; + } + NeedRecreation = true; + if (!ReconnectionCallback) { + const TDuration delay = UpdateReconnectionDelay(); + DEBUG_LOG("Schedule subproducer recreation through " << delay, Settings.SourceId, ""); + std::weak_ptr<TRetryingProducer> self(shared_from_this()); + ReconnectionCallback = + PQLib->GetScheduler().Schedule(delay, + this, + [self, sourceId = Settings.SourceId, logger = Logger] { + auto selfShared = self.lock(); + WRITE_LOG("Subproducer recreation callback. self is " << (selfShared ? "OK" : "nullptr") << ", stopping: " << (selfShared ? selfShared->Stopping : false), sourceId, "", TLOG_DEBUG, logger); + if (selfShared) { + selfShared->ReconnectionCallback = nullptr; + if (!selfShared->Stopping) { + selfShared->RecreateProducer(TInstant::Now() + selfShared->Settings.StartSessionTimeout); + } + } + }); + } +} + +void TRetryingProducer::ProcessStart(const NThreading::TFuture<TProducerCreateResponse>& f) noexcept { + INFO_LOG("Subproducer start response: " << f.GetValue().Response, Settings.SourceId, ""); + if (Stopping) + return; + if (NeedRecreation) + return; + if (!StartFuture.HasValue()) + return; + if (StartFuture.GetValue().Response.HasError()) { + WARN_LOG("Subproducer start error: " << f.GetValue().Response, Settings.SourceId, ""); + ScheduleRecreation(); + } else { + LastReconnectionDelay = TDuration::Zero(); + ReconnectionAttemptsDone = 0; + + // recreate on dead + DEBUG_LOG("Subscribe on subproducer death", Settings.SourceId, ""); + std::weak_ptr<TRetryingProducer> self(shared_from_this()); + PQLib->Subscribe(Producer->IsDead(), + this, + [self](const NThreading::TFuture<TError>& error) { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->OnProducerDead(error.GetValue()); + } + }); + + if (!StartPromise.HasValue()) { + StartPromise.SetValue(StartFuture.GetValue()); + } + + SendData(); + } +} + +void TRetryingProducer::OnProducerDead(const TError& error) { + WARN_LOG("Subproducer is dead: " << error, Settings.SourceId, ""); + ScheduleRecreation(); +} + +void TRetryingProducer::SendData() noexcept { + if (Stopping) + return; + ui64 maxSeqNo = StartFuture.GetValue().Response.GetInit().GetMaxSeqNo(); + std::deque<NThreading::TPromise<TProducerCommitResponse>> promises; + std::deque<TWriteData> values; + ui64 prevSeqNo = 0; + for (auto& d : ResendRequests) { + Y_VERIFY(d.SeqNo == 0 || d.SeqNo > prevSeqNo); + if (d.SeqNo != 0) { + prevSeqNo = d.SeqNo; + } + if (d.SeqNo != 0 && d.SeqNo <= maxSeqNo) { + promises.push_back(Promises.front()); + Promises.pop_front(); + values.push_back(d); + } else { + DelegateWriteAndSubscribe(d.SeqNo, std::move(d.Data)); + } + } + //Requests can be not checked - it is up to client + for (auto& d : Requests) { + DelegateWriteAndSubscribe(d.SeqNo, std::move(d.Data)); + } + ResendRequests.clear(); + Requests.clear(); + for (ui32 i = 0; i < promises.size(); ++i) { + TWriteResponse res; + res.MutableAck()->SetAlreadyWritten(true); + res.MutableAck()->SetSeqNo(values[i].SeqNo); + promises[i].SetValue(TProducerCommitResponse(values[i].SeqNo, std::move(values[i].Data), std::move(res))); + } +} + +void TRetryingProducer::Write(NThreading::TPromise<TProducerCommitResponse>& promise, const TProducerSeqNo seqNo, TData data) noexcept { + Y_VERIFY(data.IsEncoded()); + if (!StartFuture.Initialized()) { + TWriteResponse res; + res.MutableError()->SetDescription("producer is not ready"); + res.MutableError()->SetCode(NErrorCode::ERROR); + promise.SetValue(TProducerCommitResponse{seqNo, std::move(data), std::move(res)}); + return; + } + + Promises.push_back(promise); + + if (!StartFuture.HasValue() || NeedRecreation || Stopping || !Requests.empty() || !ResendRequests.empty()) { + Requests.emplace_back(seqNo, std::move(data)); + } else { + DelegateWriteAndSubscribe(seqNo, std::move(data)); + } +} + +void TRetryingProducer::DelegateWriteAndSubscribe(TProducerSeqNo seqNo, TData&& data) noexcept { + Y_VERIFY(InFlightRequests.size() == Futures.size()); + if (seqNo == 0) { + Futures.push_back(Producer->Write(data)); + } else { + Futures.push_back(Producer->Write(seqNo, data)); + } + InFlightRequests.emplace_back(seqNo, std::move(data)); + + std::weak_ptr<TRetryingProducer> self(shared_from_this()); + PQLib->Subscribe(Futures.back(), + this, + [self](const NThreading::TFuture<TProducerCommitResponse>&) { + auto selfShared = self.lock(); + if (selfShared) { + selfShared->ProcessFutures(); + } + }); +} + +NThreading::TFuture<TError> TRetryingProducer::IsDead() noexcept { + return IsDeadPromise.GetFuture(); +} + +void TRetryingProducer::ProcessFutures() noexcept { + NThreading::TPromise<TProducerCommitResponse> promise; + NThreading::TFuture<TProducerCommitResponse> future; + { + if (Stopping) { + return; + } + ++ToProcess; + while (ToProcess) { + if (Futures.empty() || !Futures.front().HasValue()) { + break; + } + + Y_VERIFY(InFlightRequests.size() == Futures.size()); + --ToProcess; + + const TProducerCommitResponse& response = Futures.front().GetValue(); + const TWriteData& writeData = InFlightRequests.front(); + Y_VERIFY(Promises.size() == Futures.size() + Requests.size() + ResendRequests.size()); + + if (NeedRecreation) { + ResendRequests.emplace_back(response.SeqNo, TData(writeData.Data)); + InFlightRequests.pop_front(); + Futures.pop_front(); + continue; + } + + if (Futures.front().GetValue().Response.HasError()) { + Y_VERIFY(!NeedRecreation); + ScheduleRecreation(); + WARN_LOG("Future response with error: " << response.Response, Settings.SourceId, ""); + ResendRequests.emplace_back(response.SeqNo, TData(writeData.Data)); + InFlightRequests.pop_front(); + Futures.pop_front(); + Producer = nullptr; + continue; + } + + promise = Promises.front(); + Promises.pop_front(); + future = Futures.front(); + InFlightRequests.pop_front(); + Futures.pop_front(); + promise.SetValue(future.GetValue()); + } + if (NeedRecreation && Futures.empty() && !ReconnectionCallback) { + // We need recreation, but scheduled recreation hasn't start producer because of nonempty future list. + DEBUG_LOG("Recreating subproducer after all futures were processed", Settings.SourceId, ""); + RecreateProducer(TInstant::Now() + Settings.StartSessionTimeout); + } + } +} + +void TRetryingProducer::Destroy(const TString& description) { + TError error; + error.SetDescription(description); + error.SetCode(NErrorCode::ERROR); + Destroy(error); +} + +void TRetryingProducer::SubscribeDestroyed() { + NThreading::TPromise<void> promise = ProducersDestroyed; + auto handler = [promise](const auto&) mutable { + promise.SetValue(); + }; + if (Producer) { + WaitExceptionOrAll(DestroyedPromise.GetFuture(), Producer->Destroyed()) + .Subscribe(handler); + } else { + DestroyedPromise.GetFuture() + .Subscribe(handler); + } + + DestroyPQLibRef(); +} + +void TRetryingProducer::Destroy(const TError& error) { + if (Stopping) { + return; + } + Stopping = true; + SubscribeDestroyed(); + const bool started = StartFuture.Initialized(); + Producer = nullptr; + if (started) { + Y_VERIFY(Promises.size() == ResendRequests.size() + Requests.size() + InFlightRequests.size()); + ResendRequests.insert(ResendRequests.begin(), InFlightRequests.begin(), InFlightRequests.end()); + ResendRequests.insert(ResendRequests.end(), Requests.begin(), Requests.end()); + for (auto& v : ResendRequests) { + TWriteResponse resp; + *resp.MutableError() = error; + Promises.front().SetValue(TProducerCommitResponse(v.SeqNo, std::move(v.Data), std::move(resp))); + Promises.pop_front(); + } + Y_VERIFY(Promises.empty()); + } + if (StartPromise.Initialized() && !StartPromise.HasValue()) { + TWriteResponse resp; + *resp.MutableError() = error; + StartPromise.SetValue(TProducerCreateResponse(std::move(resp))); + } + IsDeadPromise.SetValue(error); +} + +NThreading::TFuture<void> TRetryingProducer::Destroyed() noexcept { + return ProducersDestroyed.GetFuture(); +} + +void TRetryingProducer::OnStartDeadline() { + if (!StartPromise.HasValue()) { + TError error; + error.SetDescription("Start timeout."); + error.SetCode(NErrorCode::CREATE_TIMEOUT); + Destroy(error); + } +} + +void TRetryingProducer::Cancel() { + Destroy(GetCancelReason()); +} + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_producer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_producer.h new file mode 100644 index 0000000000..f1242564b9 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_producer.h @@ -0,0 +1,75 @@ +#pragma once + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include "scheduler.h" +#include "internals.h" +#include "persqueue_p.h" +#include "producer.h" +#include "iproducer_p.h" +#include <library/cpp/threading/future/future.h> +#include <kikimr/yndx/api/grpc/persqueue.grpc.pb.h> +#include <deque> + +namespace NPersQueue { + +class TRetryingProducer: public IProducerImpl, public std::enable_shared_from_this<TRetryingProducer> { +public: + using IProducerImpl::Write; + // If seqno == 0, we assume it to be autodefined + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, const TProducerSeqNo seqNo, TData data) noexcept override; + void Write(NThreading::TPromise<TProducerCommitResponse>& promise, TData data) noexcept override { + TRetryingProducer::Write(promise, 0, std::move(data)); + } + + TRetryingProducer(const TProducerSettings& settings, std::shared_ptr<void> destroyEventRef, TIntrusivePtr<TPQLibPrivate> pqLib, + TIntrusivePtr<ILogger> logger); + + NThreading::TFuture<TProducerCreateResponse> Start(TInstant deadline) noexcept override; + + NThreading::TFuture<TError> IsDead() noexcept override; + + NThreading::TFuture<void> Destroyed() noexcept override; + + ~TRetryingProducer() noexcept; + + + void RecreateProducer(TInstant deadline) noexcept; + void ProcessStart(const NThreading::TFuture<TProducerCreateResponse>&) noexcept; + void SendData() noexcept; + void ProcessFutures() noexcept; + void DoRecreate(TInstant deadline) noexcept; + void DelegateWriteAndSubscribe(TProducerSeqNo seqNo, TData&& data) noexcept; + TDuration UpdateReconnectionDelay(); + void ScheduleRecreation(); + void Destroy(const TError& error); + void Destroy(const TString& description); // the same but with Code=ERROR + void SubscribeDestroyed(); + void OnStartDeadline(); + void OnProducerDead(const TError& error); + + void Cancel() override; + +protected: + TProducerSettings Settings; + TIntrusivePtr<ILogger> Logger; + std::shared_ptr<IProducerImpl> Producer; + std::deque<TWriteData> Requests; + std::deque<TWriteData> ResendRequests; + std::deque<TWriteData> InFlightRequests; + std::deque<NThreading::TPromise<TProducerCommitResponse>> Promises; + std::deque<NThreading::TFuture<TProducerCommitResponse>> Futures; + NThreading::TFuture<TProducerCreateResponse> StartFuture; + NThreading::TPromise<TProducerCreateResponse> StartPromise; + NThreading::TPromise<TError> IsDeadPromise; + bool NeedRecreation; + bool Stopping; + ui32 ToProcess; + TDuration LastReconnectionDelay; + unsigned ReconnectionAttemptsDone; + TIntrusivePtr<TScheduler::TCallbackHandler> ReconnectionCallback; + NThreading::TPromise<void> ProducersDestroyed = NThreading::NewPromise<void>(); +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_producer_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_producer_ut.cpp new file mode 100644 index 0000000000..3fdb17d954 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/retrying_producer_ut.cpp @@ -0,0 +1,355 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h> + +#include <library/cpp/testing/unittest/registar.h> + +using namespace NThreading; +using namespace NKikimr; +using namespace NKikimr::NPersQueueTests; + +namespace NPersQueue { + +Y_UNIT_TEST_SUITE(TRetryingProducerTest) { + TProducerSettings FakeSettings() { + TProducerSettings settings; + settings.ReconnectOnFailure = true; + settings.Server = TServerSetting{"localhost"}; + settings.Topic = "topic"; + settings.SourceId = "src"; + return settings; + } + + Y_UNIT_TEST(NotStartedProducerCanBeDestructed) { + // Test that producer doesn't hang on till shutdown + TPQLib lib; + TProducerSettings settings = FakeSettings(); + lib.CreateProducer(settings, {}, false); + } + + Y_UNIT_TEST(StartDeadlineExpires) { + if (std::getenv("PERSQUEUE_GRPC_API_V1_ENABLED")) + return; + TPQLibSettings libSettings; + libSettings.ChannelCreationTimeout = TDuration::MilliSeconds(1); + libSettings.DefaultLogger = new TCerrLogger(TLOG_DEBUG); + TPQLib lib(libSettings); + THolder<IProducer> producer; + { + TProducerSettings settings = FakeSettings(); + settings.ReconnectionDelay = TDuration::MilliSeconds(1); + settings.StartSessionTimeout = TDuration::MilliSeconds(10); + producer = lib.CreateProducer(settings, {}, false); + } + const TInstant beforeStart = TInstant::Now(); + auto future = producer->Start(TDuration::MilliSeconds(100)); + UNIT_ASSERT(future.GetValueSync().Response.HasError()); + UNIT_ASSERT_EQUAL_C(future.GetValueSync().Response.GetError().GetCode(), NErrorCode::CREATE_TIMEOUT, + "Error: " << future.GetValueSync().Response.GetError()); + const TInstant now = TInstant::Now(); + UNIT_ASSERT_C(now - beforeStart >= TDuration::MilliSeconds(100), now); + + DestroyAndWait(producer); + } + + Y_UNIT_TEST(StartMaxAttemptsExpire) { + TPQLibSettings libSettings; + libSettings.ChannelCreationTimeout = TDuration::MilliSeconds(1); + TPQLib lib(libSettings); + THolder<IProducer> producer; + { + TProducerSettings settings = FakeSettings(); + settings.ReconnectionDelay = TDuration::MilliSeconds(10); + settings.StartSessionTimeout = TDuration::MilliSeconds(10); + settings.MaxAttempts = 3; + producer = lib.CreateProducer(settings, {}, false); + } + const TInstant beforeStart = TInstant::Now(); + auto future = producer->Start(); + UNIT_ASSERT(future.GetValueSync().Response.HasError()); + const TInstant now = TInstant::Now(); + UNIT_ASSERT_C(now - beforeStart >= TDuration::MilliSeconds(30), now); + + DestroyAndWait(producer); + } + + void AssertWriteValid(const NThreading::TFuture<TProducerCommitResponse>& respFuture) { + const TProducerCommitResponse& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TWriteResponse::kAck, "Msg: " << resp.Response); + } + + void AssertWriteFailed(const NThreading::TFuture<TProducerCommitResponse>& respFuture) { + const TProducerCommitResponse& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TWriteResponse::kError, "Msg: " << resp.Response); + } + + void AssertReadValid(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kData, "Msg: " << resp.Response); + } + + void AssertReadContinuous(const TReadResponse::TData::TMessageBatch& batch, ui64 startSeqNo) { + for (ui32 i = 0; i < batch.MessageSize(); ++i) { + ui64 actualSeqNo = batch.GetMessage(i).GetMeta().GetSeqNo(); + UNIT_ASSERT_EQUAL_C(actualSeqNo, startSeqNo + i, "Wrong seqNo: " << actualSeqNo << ", expected: " << startSeqNo + i); + } + } + + void AssertCommited(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kCommit, "Msg: " << resp.Response); + } + + void AssertReadFailed(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kError, "Msg: " << resp.Response); + } + + void AssertLock(const NThreading::TFuture<TConsumerMessage>& respFuture) { + const TConsumerMessage& resp = respFuture.GetValueSync(); + UNIT_ASSERT_EQUAL_C(resp.Response.GetResponseCase(), TReadResponse::kLock, "Msg: " << resp.Response); + } + + TProducerSettings MakeProducerSettings(const TTestServer& testServer) { + TProducerSettings producerSettings; + producerSettings.ReconnectOnFailure = true; + producerSettings.Topic = "topic1"; + producerSettings.SourceId = "123"; + producerSettings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + producerSettings.Codec = ECodec::LZOP; + producerSettings.ReconnectionDelay = TDuration::MilliSeconds(10); + return producerSettings; + } + + TConsumerSettings MakeConsumerSettings(const TTestServer& testServer) { + TConsumerSettings consumerSettings; + consumerSettings.ClientId = "user"; + consumerSettings.Server = TServerSetting{"localhost", testServer.GrpcPort}; + consumerSettings.Topics.push_back("topic1"); + return consumerSettings; + } + + Y_UNIT_TEST(ReconnectsToServer) { + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(false); + + testServer.AnnoyingClient->FullInit(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + const size_t partitions = 10; + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", partitions); + testServer.WaitInit("topic1"); + + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLib PQLib; + + auto producer = PQLib.CreateProducer(MakeProducerSettings(testServer), logger, false); + producer->Start().Wait(); + TFuture<TError> isDead = producer->IsDead(); + UNIT_ASSERT(!isDead.HasValue()); + + auto consumer = PQLib.CreateConsumer(MakeConsumerSettings(testServer), logger, false); + consumer->Start().Wait(); + + auto read1 = consumer->GetNextMessage(); + + // write first value + auto write1 = producer->Write(1, TString("blob1")); + AssertWriteValid(write1); + AssertReadValid(read1); + + // commit + consumer->Commit({1}); + auto commitAck = consumer->GetNextMessage(); + AssertCommited(commitAck); + + auto read2 = consumer->GetNextMessage(); + testServer.ShutdownServer(); + UNIT_ASSERT(!isDead.HasValue()); + + auto write2 = producer->Write(2, TString("blob2")); + + testServer.StartServer(false); + + testServer.AnnoyingClient->FullInit(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", partitions); + + AssertWriteValid(write2); + + AssertReadFailed(read2); + consumer->IsDead().Wait(); + + consumer = PQLib.CreateConsumer(MakeConsumerSettings(testServer), logger, false); + consumer->Start().Wait(); + + read2 = consumer->GetNextMessage(); + AssertReadValid(read2); + + UNIT_ASSERT(!isDead.HasValue()); + + DestroyAndWait(producer); + DestroyAndWait(consumer); + } + + static void DiesOnTooManyReconnectionAttempts(bool callWrite) { + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(false); + + testServer.AnnoyingClient->FullInit(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", 2); + testServer.WaitInit("topic1"); + + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLib PQLib; + + TProducerSettings settings = MakeProducerSettings(testServer); + settings.MaxAttempts = 3; + settings.ReconnectionDelay = TDuration::MilliSeconds(100); + auto producer = PQLib.CreateProducer(settings, logger, false); + producer->Start().Wait(); + TFuture<TError> isDead = producer->IsDead(); + UNIT_ASSERT(!isDead.HasValue()); + + // shutdown server + const TInstant beforeShutdown = TInstant::Now(); + testServer.ShutdownServer(); + + NThreading::TFuture<TProducerCommitResponse> write; + if (callWrite) { + write = producer->Write(TString("data")); + } + + isDead.Wait(); + const TInstant afterDead = TInstant::Now(); + // 3 attempts: 100ms, 200ms and 400ms + UNIT_ASSERT_C(afterDead - beforeShutdown >= TDuration::MilliSeconds(700), "real difference: " << (afterDead - beforeShutdown)); + + if (callWrite) { + AssertWriteFailed(write); + } + } + + Y_UNIT_TEST(WriteDataThatWasSentBeforeConnectToServer) { + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(false); + + testServer.AnnoyingClient->FullInit(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + const size_t partitions = 10; + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", partitions); + + testServer.WaitInit("topic1"); + + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLib PQLib; + + auto producer = PQLib.CreateProducer(MakeProducerSettings(testServer), logger, false); + auto producerStarted = producer->Start(); + auto write1 = producer->Write(1, TString("a")); + producerStarted.Wait(); + TVector<TFuture<TProducerCommitResponse>> acks; + for (int i = 2; i <= 1000; ++i) { + acks.push_back(producer->Write(i, TString("b"))); + } + + AssertWriteValid(write1); + for (auto& ack: acks) { + ack.Wait(); + } + + auto consumer = PQLib.CreateConsumer(MakeConsumerSettings(testServer), logger, false); + consumer->Start().Wait(); + + ui64 readMessages = 0; + while (readMessages < 1000) { + auto read1 = consumer->GetNextMessage(); + AssertReadValid(read1); + auto& batch = read1.GetValueSync().Response.GetData().GetMessageBatch(0); + AssertReadContinuous(batch, readMessages + 1); + readMessages += batch.MessageSize(); + } + + DestroyAndWait(producer); + DestroyAndWait(consumer); + } + + Y_UNIT_TEST(DiesOnTooManyReconnectionAttemptsWithoutWrite) { + // Check that we reconnect even without explicit write errors + DiesOnTooManyReconnectionAttempts(false); + } + + Y_UNIT_TEST(DiesOnTooManyReconnectionAttemptsWithWrite) { + DiesOnTooManyReconnectionAttempts(true); + } + + Y_UNIT_TEST(CancelsOperationsAfterPQLibDeath) { + return; // Test is ignored. FIX: KIKIMR-7886 + TTestServer testServer(false); + testServer.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::MilliSeconds(10)); + testServer.StartServer(); + + const size_t partitions = 1; + testServer.AnnoyingClient->InitRoot(); + testServer.AnnoyingClient->InitDCs(!GrpcV1EnabledByDefault() ? DEFAULT_CLUSTERS_LIST : CLUSTERS_LIST_ONE_DC); + testServer.AnnoyingClient->CreateTopicNoLegacy("rt3.dc1--topic1", partitions); + testServer.AnnoyingClient->InitSourceIds(); + + testServer.WaitInit("topic1"); + + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLibSettings pqLibSettings; + pqLibSettings.DefaultLogger = logger; + THolder<TPQLib> PQLib = MakeHolder<TPQLib>(pqLibSettings); + + auto producer = PQLib->CreateProducer(MakeProducerSettings(testServer), logger, false); + UNIT_ASSERT(!producer->Start().GetValueSync().Response.HasError()); + TFuture<TError> isDead = producer->IsDead(); + UNIT_ASSERT(!isDead.HasValue()); + + testServer.ShutdownServer(); + UNIT_ASSERT(!isDead.HasValue()); + + auto write1 = producer->Write(1, TString("blob1")); + auto write2 = producer->Write(2, TString("blob2")); + + UNIT_ASSERT(!write1.HasValue()); + UNIT_ASSERT(!write2.HasValue()); + + PQLib = nullptr; + + UNIT_ASSERT(write1.HasValue()); + UNIT_ASSERT(write2.HasValue()); + + UNIT_ASSERT(write1.GetValue().Response.HasError()); + UNIT_ASSERT(write2.GetValue().Response.HasError()); + + auto write3 = producer->Write(3, TString("blob3")); + UNIT_ASSERT(write3.HasValue()); + UNIT_ASSERT(write3.GetValue().Response.HasError()); + } + + Y_UNIT_TEST(CancelsStartAfterPQLibDeath) { + TIntrusivePtr<TCerrLogger> logger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + TPQLibSettings pqLibSettings; + pqLibSettings.DefaultLogger = logger; + THolder<TPQLib> PQLib = MakeHolder<TPQLib>(pqLibSettings); + + auto producer = PQLib->CreateProducer(FakeSettings(), logger, false); + auto start = producer->Start(); + UNIT_ASSERT(!start.HasValue()); + + PQLib = nullptr; + + UNIT_ASSERT(start.HasValue()); + UNIT_ASSERT(start.GetValue().Response.HasError()); + + auto dead = producer->IsDead(); + UNIT_ASSERT(dead.HasValue()); + dead.GetValueSync(); + + auto write = producer->Write(1, TString("blob1")); + UNIT_ASSERT(write.HasValue()); + UNIT_ASSERT(write.GetValueSync().Response.HasError()); + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/scheduler.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/scheduler.cpp new file mode 100644 index 0000000000..ba6c4ebb40 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/scheduler.cpp @@ -0,0 +1,96 @@ +#include "persqueue_p.h" +#include "scheduler.h" + +#include <util/system/thread.h> + +namespace NPersQueue { +void TScheduler::TCallbackHandler::Execute() { + auto guard = Guard(Lock); + if (Callback) { + Y_VERIFY(PQLib->GetQueuePool().GetQueue(QueueTag).AddFunc(std::move(Callback))); + Callback = nullptr; + } +} + +void TScheduler::TCallbackHandler::TryCancel() { + auto guard = Guard(Lock); + if (Callback) { + Callback = nullptr; + } +} + +bool TScheduler::TCallbackHandlersCompare::operator()(const TIntrusivePtr<TScheduler::TCallbackHandler>& h1, const TIntrusivePtr<TScheduler::TCallbackHandler>& h2) const { + return h1->Time > h2->Time; +} + +TScheduler::TScheduler(TPQLibPrivate* pqLib) + : Shutdown(false) + , PQLib(pqLib) +{ + Thread = SystemThreadFactory()->Run([this] { + this->SchedulerThread(); + }); +} + +TScheduler::~TScheduler() { + ShutdownAndWait(); +} + +void TScheduler::AddToSchedule(TIntrusivePtr<TCallbackHandler> handler) { + { + auto guard = Guard(Lock); + Callbacks.push(std::move(handler)); + } + Event.Signal(); +} + +void TScheduler::ShutdownAndWait() { + AtomicSet(Shutdown, true); + Event.Signal(); + Thread->Join(); +} + +void TScheduler::SchedulerThread() { + TThread::SetCurrentThreadName("pqlib_scheduler"); + while (!AtomicGet(Shutdown)) { + TInstant deadline = TInstant::Max(); + std::vector<TIntrusivePtr<TCallbackHandler>> callbacks; + { + // define next deadline and get expired callbacks + auto guard = Guard(Lock); + const TInstant now = TInstant::Now(); + while (!Callbacks.empty()) { + const auto& top = Callbacks.top(); + if (top->Time <= now) { + callbacks.push_back(top); + Callbacks.pop(); + } else { + deadline = top->Time; + break; + } + } + } + + // execute callbacks + bool shutdown = false; + for (auto& callback : callbacks) { + if (shutdown) { + callback->TryCancel(); + } else { + callback->Execute(); + } + shutdown = shutdown || AtomicGet(Shutdown); + } + + if (!shutdown) { + Event.WaitD(deadline); + } + } + + // cancel all callbacks and clear data + while (!Callbacks.empty()) { + Callbacks.top()->TryCancel(); + Callbacks.pop(); + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/scheduler.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/scheduler.h new file mode 100644 index 0000000000..c90ab87795 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/scheduler.h @@ -0,0 +1,88 @@ +#pragma once +#include <util/datetime/base.h> +#include <util/generic/ptr.h> +#include <util/system/event.h> +#include <util/system/spinlock.h> +#include <util/thread/factory.h> + +#include <functional> +#include <queue> +#include <utility> + +namespace NPersQueue { + +class TPQLibPrivate; + +class TScheduler { + struct TCallbackHandlersCompare; +public: + using TCallback = std::function<void()>; + + class TCallbackHandler: public TAtomicRefCount<TCallbackHandler> { + friend class TScheduler; + friend struct TCallbackHandlersCompare; + + template <class TFunc> + TCallbackHandler(TInstant time, TFunc&& func, TPQLibPrivate* pqLib, const void* queueTag) + : Time(time) + , Callback(std::forward<TFunc>(func)) + , PQLib(pqLib) + , QueueTag(queueTag) + { + } + + void Execute(); + + public: + // Cancels execution of callback. + // If callback is already canceled or executes (executed), does nothing. + // Posteffect: callback is guaranteed to be destroyed after this call. + void TryCancel(); + + private: + TInstant Time; + TCallback Callback; + TAdaptiveLock Lock; + TPQLibPrivate* PQLib; + const void* QueueTag; + }; + +public: + // Starts a scheduler thread. + explicit TScheduler(TPQLibPrivate* pqLib); + + // Stops a scheduler thread. + ~TScheduler(); + + // Schedules a new callback to be executed + template <class TFunc> + TIntrusivePtr<TCallbackHandler> Schedule(TInstant time, const void* queueTag, TFunc&& func) { + TIntrusivePtr<TCallbackHandler> handler(new TCallbackHandler(time, std::forward<TFunc>(func), PQLib, queueTag)); + AddToSchedule(handler); + return handler; + } + + template <class TFunc> + TIntrusivePtr<TCallbackHandler> Schedule(TDuration delta, const void* queueTag, TFunc&& func) { + return Schedule(TInstant::Now() + delta, queueTag, std::forward<TFunc>(func)); + } + +private: + void AddToSchedule(TIntrusivePtr<TCallbackHandler> handler); + void ShutdownAndWait(); + void SchedulerThread(); + +private: + struct TCallbackHandlersCompare { + bool operator()(const TIntrusivePtr<TCallbackHandler>& h1, const TIntrusivePtr<TCallbackHandler>& h2) const; + }; + + using TCallbackQueue = std::priority_queue<TIntrusivePtr<TCallbackHandler>, std::vector<TIntrusivePtr<TCallbackHandler>>, TCallbackHandlersCompare>; + TAdaptiveLock Lock; + TAutoEvent Event; + TCallbackQueue Callbacks; + TAtomic Shutdown; + THolder<IThreadFactory::IThread> Thread; + TPQLibPrivate* PQLib; +}; +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/scheduler_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/scheduler_ut.cpp new file mode 100644 index 0000000000..83ce25c6b4 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/scheduler_ut.cpp @@ -0,0 +1,75 @@ +#include "scheduler.h" +#include "persqueue_p.h" + +#include <library/cpp/testing/unittest/registar.h> + +#include <util/system/event.h> + +namespace NPersQueue { +const void* Tag = (const void*)42; +Y_UNIT_TEST_SUITE(TSchedulerTest) { + Y_UNIT_TEST(EmptySchedulerCanDestruct) { + // Test that scheduler doesn't hang on till shutdown + TScheduler scheduler(nullptr); + } + + Y_UNIT_TEST(ExecutesInProperOrder) { + TSystemEvent ev; + bool callback1Executed = false; + TIntrusivePtr<TPQLibPrivate> lib(new TPQLibPrivate({})); + TScheduler scheduler(lib.Get()); + scheduler.Schedule(TDuration::MilliSeconds(100), Tag, [&callback1Executed, &ev] { + UNIT_ASSERT(callback1Executed); + ev.Signal(); + }); + scheduler.Schedule(TDuration::MilliSeconds(50), Tag, [&callback1Executed] { + callback1Executed = true; + }); + ev.Wait(); + } + + Y_UNIT_TEST(CancelsAndClearsData) { + TIntrusivePtr<TPQLibPrivate> lib(new TPQLibPrivate({})); + TScheduler scheduler(lib.Get()); + std::shared_ptr<TString> ptr(new TString()); + bool callbackExecuted = false; + auto h = scheduler.Schedule(TDuration::Seconds(50), Tag, [&callbackExecuted, ptr] { + callbackExecuted = true; + }); + UNIT_ASSERT_VALUES_EQUAL(ptr.use_count(), 2); + h->TryCancel(); + UNIT_ASSERT_VALUES_EQUAL(ptr.use_count(), 1); + UNIT_ASSERT(!callbackExecuted); + Sleep(TDuration::MilliSeconds(51)); + UNIT_ASSERT_VALUES_EQUAL(ptr.use_count(), 1); + UNIT_ASSERT(!callbackExecuted); + } + + Y_UNIT_TEST(ExitsThreadImmediately) { + TIntrusivePtr<TPQLibPrivate> lib(new TPQLibPrivate({})); + std::shared_ptr<TString> ptr(new TString()); + bool callback1Executed = false; + bool callback2Executed = false; + TSystemEvent ev, ev2; + auto now = TInstant::Now(); + lib->GetScheduler().Schedule(TDuration::Seconds(500), Tag, [&callback1Executed, ptr] { + callback1Executed = true; + }); + lib->GetScheduler().Schedule(now, Tag, [&callback2Executed, &ev, &ev2, ptr] { + callback2Executed = true; + ev2.Wait(); + ev.Signal(); + }); + UNIT_ASSERT_VALUES_EQUAL(ptr.use_count(), 3); + ev2.Signal(); + + // kill scheduler + ev.Wait(); + lib.Reset(); + + UNIT_ASSERT_VALUES_EQUAL(ptr.use_count(), 1); + UNIT_ASSERT(!callback1Executed); + UNIT_ASSERT(callback2Executed); + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/types.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/types.cpp new file mode 100644 index 0000000000..4fc6a1c7b4 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/types.cpp @@ -0,0 +1,86 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> + +#include <library/cpp/streams/lzop/lzop.h> +#include <library/cpp/streams/zstd/zstd.h> + +#include <util/generic/store_policy.h> +#include <util/generic/utility.h> +#include <util/stream/str.h> +#include <util/stream/zlib.h> + +namespace NPersQueue { + +const TString& TServerSetting::GetRootDatabase() { + static const TString RootDatabase = "/Root"; + return RootDatabase; +} + +class TZLibToStringCompressor: private TEmbedPolicy<TStringOutput>, public TZLibCompress { +public: + TZLibToStringCompressor(TString& dst, ZLib::StreamType type, size_t quality) + : TEmbedPolicy<TStringOutput>(dst) + , TZLibCompress(TEmbedPolicy::Ptr(), type, quality) + { + } +}; + +class TLzopToStringCompressor: private TEmbedPolicy<TStringOutput>, public TLzopCompress { +public: + TLzopToStringCompressor(TString& dst) + : TEmbedPolicy<TStringOutput>(dst) + , TLzopCompress(TEmbedPolicy::Ptr()) + { + } +}; + +class TZstdToStringCompressor: private TEmbedPolicy<TStringOutput>, public TZstdCompress { +public: + TZstdToStringCompressor(TString& dst, int quality) + : TEmbedPolicy<TStringOutput>(dst) + , TZstdCompress(TEmbedPolicy::Ptr(), quality) + { + } +}; + + +TData TData::Encode(TData source, ECodec defaultCodec, int quality) { + Y_VERIFY(!source.Empty()); + TData data = std::move(source); + if (data.IsEncoded()) { + return data; + } + Y_VERIFY(defaultCodec != ECodec::RAW && defaultCodec != ECodec::DEFAULT); + if (data.Codec == ECodec::DEFAULT) { + data.Codec = defaultCodec; + } + THolder<IOutputStream> coder = CreateCoder(data.Codec, data, quality); + coder->Write(data.SourceData); + coder->Finish(); // &data.EncodedData may be already invalid on coder destruction + return data; +} + +TData TData::MakeRawIfNotEncoded(TData source) { + Y_VERIFY(!source.Empty()); + TData data = std::move(source); + if (!data.IsEncoded()) { + data.EncodedData = data.SourceData; + data.Codec = ECodec::RAW; + } + return data; +} + +THolder<IOutputStream> TData::CreateCoder(ECodec codec, TData& result, int quality) { + switch (codec) { + case ECodec::GZIP: + return MakeHolder<TZLibToStringCompressor>(result.EncodedData, ZLib::GZip, quality >= 0 ? quality : 6); + case ECodec::LZOP: + return MakeHolder<TLzopToStringCompressor>(result.EncodedData); + case ECodec::ZSTD: + return MakeHolder<TZstdToStringCompressor>(result.EncodedData, quality); + default: + Y_FAIL("NOT IMPLEMENTED CODEC TYPE"); + } +} + + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/types_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/types_ut.cpp new file mode 100644 index 0000000000..06c4b28454 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/types_ut.cpp @@ -0,0 +1,88 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> + +#include <library/cpp/testing/unittest/registar.h> + +#include <util/string/builder.h> + +namespace NPersQueue { +Y_UNIT_TEST_SUITE(TDataTest) { + ECodec Codecs[] = { + ECodec::LZOP, + ECodec::GZIP, + }; + + TString DebugString(const TData& data) { + return TStringBuilder() << "{ Ts: " << data.GetTimestamp().GetValue() << ", SrcData: \"" << data.GetSourceData() + << "\", Encoded: \"" << (data.IsEncoded() ? data.GetEncodedData() : TString()) + << "\", Codec: " << data.GetCodecType() << " }"; + } + + Y_UNIT_TEST(RawDataIsEncoded) { + auto now = TInstant::Now(); + TData data = TData::Raw("trololo", now); + TData otherRawData = TData("trololo", ECodec::RAW, now); + UNIT_ASSERT_C(data.IsEncoded(), "data: " << DebugString(data)); + UNIT_ASSERT_C(!data.Empty(), "data: " << DebugString(data)); + UNIT_ASSERT_EQUAL_C(data.GetCodecType(), ECodec::RAW, "data: " << DebugString(data)); + UNIT_ASSERT_EQUAL_C(data, otherRawData, "data: " << DebugString(data) << ", other: " << DebugString(otherRawData)); + } + + Y_UNIT_TEST(EncodedDataIsEncoded) { + for (ECodec codec : Codecs) { + TData data = TData::Encoded("trololo", codec); + UNIT_ASSERT(data.IsEncoded()); + UNIT_ASSERT(!data.Empty()); + UNIT_ASSERT_EQUAL(data.GetCodecType(), codec); + } + } + + Y_UNIT_TEST(ModifiesState) { + for (ECodec codec : Codecs) { + for (ECodec defaultCodec : Codecs) { + TData data("trololo", codec); + UNIT_ASSERT(!data.IsEncoded()); + UNIT_ASSERT(!data.Empty()); + + for (size_t i = 0; i < 2; ++i) { + data = TData::Encode(std::move(data), defaultCodec, 2); // encode twice is OK + UNIT_ASSERT(data.IsEncoded()); + UNIT_ASSERT(!data.Empty()); + UNIT_ASSERT(!data.GetEncodedData().empty()); + UNIT_ASSERT_EQUAL(data.GetCodecType(), codec); + } + } + } + } + + Y_UNIT_TEST(HandlesDefaultCodec) { + for (ECodec defaultCodec : Codecs) { + TData data = TString("trololo"); + UNIT_ASSERT(!data.IsEncoded()); + UNIT_ASSERT(!data.Empty()); + + data = TData::Encode(data, defaultCodec, -1); + UNIT_ASSERT(data.IsEncoded()); + UNIT_ASSERT(!data.Empty()); + UNIT_ASSERT_STRINGS_EQUAL(data.GetSourceData(), "trololo"); + UNIT_ASSERT(!data.GetEncodedData().empty()); + UNIT_ASSERT_EQUAL(data.GetCodecType(), defaultCodec); + } + } + + Y_UNIT_TEST(MakesRaw) { + TData data = TString("trololo"); + UNIT_ASSERT(!data.IsEncoded()); + data = TData::MakeRawIfNotEncoded(data); + UNIT_ASSERT(data.IsEncoded()); + UNIT_ASSERT_EQUAL(data.GetCodecType(), ECodec::RAW); + } + + Y_UNIT_TEST(DoesNotMakeRaw) { + TData data = TData::Encoded("trololo", ECodec::GZIP); + UNIT_ASSERT(data.IsEncoded()); + data = TData::MakeRawIfNotEncoded(data); + UNIT_ASSERT(data.IsEncoded()); + UNIT_ASSERT_EQUAL(data.GetCodecType(), ECodec::GZIP); + } +} +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/validate_grpc_metadata.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/validate_grpc_metadata.cpp new file mode 100644 index 0000000000..a7d65ddc88 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/validate_grpc_metadata.cpp @@ -0,0 +1,24 @@ +#include "validate_grpc_metadata.h" +#include <contrib/libs/grpc/include/grpc/grpc.h> +#include <util/generic/string.h> +#include <util/string/escape.h> +#include <util/string/builder.h> + +namespace NGrpc { + bool ValidateHeaderIsLegal(const TString& key, const TString& value, TString& error) { + error.clear(); + grpc_slice keySlice = grpc_slice_from_static_buffer(key.c_str(), key.size()); + int ok = grpc_header_key_is_legal(keySlice); + if (!ok) { + error = TStringBuilder() << "gRPC metadata header key is illegal: \"" << EscapeC(key) << "\""; + return false; + } + grpc_slice valueSlice = grpc_slice_from_static_buffer(value.c_str(), value.size()); + ok = grpc_is_binary_header(keySlice) || grpc_header_nonbin_value_is_legal(valueSlice); + if (!ok) { + error = TStringBuilder() << "gRPC metadata header value with key \"" << key << "\" is illegal: \"" << EscapeC(value) << "\""; + return false; + } + return true; + } +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/validate_grpc_metadata.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/validate_grpc_metadata.h new file mode 100644 index 0000000000..43fd82cef0 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/validate_grpc_metadata.h @@ -0,0 +1,12 @@ +#pragma once +#include <util/generic/fwd.h> + +namespace NGrpc { + // Validates gRPC metadata header key and value. Returns 'true' if validations fails, otherwise returns 'false' and sets 'error' value. + // For more information see 'Custom-Metadata' section in gRPC over HTTP2 (https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md). + // Note that in case of authentication data validation 'error' may contain sensitive information. + bool ValidateHeaderIsLegal(const TString& key, const TString& value, TString& error); + + // Creates hexadecimal representation of 's' in format '{0x01, 0x02, 0x03}' + TString ToHexString(const TString& s); +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/validate_grpc_metadata_ut.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/validate_grpc_metadata_ut.cpp new file mode 100644 index 0000000000..bbda57f538 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/validate_grpc_metadata_ut.cpp @@ -0,0 +1,46 @@ +#include "validate_grpc_metadata.h" +#include <library/cpp/testing/unittest/registar.h> + +using namespace NGrpc; + +Y_UNIT_TEST_SUITE(NValidateGrpcMetadata) { + Y_UNIT_TEST(ValidateHeaderIsLegal) { + TString error = "error"; + UNIT_ASSERT_C(ValidateHeaderIsLegal("key", "value", error), error); + // Assert 'error' is cleared + UNIT_ASSERT_C(error.empty(), error); + + // Valid character values upper and lower bounds + UNIT_ASSERT_C(ValidateHeaderIsLegal("\x30\x39\x61\x7A_-.", "\x20\x7E", error), error); + UNIT_ASSERT_C(error.empty(), error); + + TString null = " "; + null[0] = '\0'; + UNIT_ASSERT(!ValidateHeaderIsLegal("key", null, error)); + UNIT_ASSERT_C(error.Contains("\\0"), error); + Cerr << "Error is '" << error << "'" << Endl; + + // Simple escape sequences + UNIT_ASSERT(!ValidateHeaderIsLegal("key", "value\x0A\t", error)); + UNIT_ASSERT_C(error.Contains("\\n\\t"), error); + Cerr << "Error is '" << error << "'" << Endl; + + UNIT_ASSERT(!ValidateHeaderIsLegal("key", "value\x1F", error)); + UNIT_ASSERT_C(error.Contains("\\x1F"), error); + Cerr << "Error is '" << error << "'" << Endl; + + UNIT_ASSERT(!ValidateHeaderIsLegal("key", "value\x7F", error)); + UNIT_ASSERT_C(error.Contains("\\x7F"), error); + Cerr << "Error is '" << error << "'" << Endl; + + // Octal character + UNIT_ASSERT(!ValidateHeaderIsLegal("key", "value\177", error)); + UNIT_ASSERT_C(error.Contains("\\x7F"), error); + Cerr << "Error is '" << error << "'" << Endl; + + // Invalid header key + UNIT_ASSERT(!ValidateHeaderIsLegal("key\n", "value", error)); + UNIT_ASSERT_C(error.Contains("\\n"), error); + Cerr << "Error is '" << error << "'" << Endl; + } +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/ydb_sdk_consumer.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/ydb_sdk_consumer.cpp new file mode 100644 index 0000000000..b679eda9f4 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/ydb_sdk_consumer.cpp @@ -0,0 +1,561 @@ +#include "ydb_sdk_consumer.h" +#include "persqueue_p.h" + +#include <ydb/library/persqueue/topic_parser_public/topic_parser.h> + +#include <ydb/public/sdk/cpp/client/ydb_persqueue_core/impl/common.h> +#include <util/generic/is_in.h> +#include <util/string/builder.h> + +namespace NPersQueue { + +static TError MakeError(const TString& description, NErrorCode::EErrorCode code) { + TError error; + error.SetDescription(description); + error.SetCode(code); + return error; +} + +static TString MakeLegacyTopicName(const NYdb::NPersQueue::TPartitionStream::TPtr& partitionStream) { + // rt3.man--account--topic + return BuildFullTopicName(partitionStream->GetTopicPath(), partitionStream->GetCluster()); +} + +static NErrorCode::EErrorCode ToErrorCode(NYdb::EStatus status) { + switch (status) { + case NYdb::EStatus::STATUS_UNDEFINED: + return NErrorCode::ERROR; + case NYdb::EStatus::SUCCESS: + return NErrorCode::OK; + case NYdb::EStatus::BAD_REQUEST: + return NErrorCode::BAD_REQUEST; + case NYdb::EStatus::UNAUTHORIZED: + return NErrorCode::ACCESS_DENIED; + case NYdb::EStatus::INTERNAL_ERROR: + return NErrorCode::ERROR; + case NYdb::EStatus::ABORTED: + [[fallthrough]]; + case NYdb::EStatus::UNAVAILABLE: + return NErrorCode::ERROR; + case NYdb::EStatus::OVERLOADED: + return NErrorCode::OVERLOAD; + case NYdb::EStatus::SCHEME_ERROR: + return NErrorCode::UNKNOWN_TOPIC; + case NYdb::EStatus::GENERIC_ERROR: + [[fallthrough]]; + case NYdb::EStatus::TIMEOUT: + [[fallthrough]]; + case NYdb::EStatus::BAD_SESSION: + [[fallthrough]]; + case NYdb::EStatus::PRECONDITION_FAILED: + [[fallthrough]]; + case NYdb::EStatus::ALREADY_EXISTS: + [[fallthrough]]; + case NYdb::EStatus::NOT_FOUND: + [[fallthrough]]; + case NYdb::EStatus::SESSION_EXPIRED: + [[fallthrough]]; + case NYdb::EStatus::CANCELLED: + [[fallthrough]]; + case NYdb::EStatus::UNDETERMINED: + [[fallthrough]]; + case NYdb::EStatus::UNSUPPORTED: + [[fallthrough]]; + case NYdb::EStatus::SESSION_BUSY: + [[fallthrough]]; + case NYdb::EStatus::TRANSPORT_UNAVAILABLE: + [[fallthrough]]; + case NYdb::EStatus::CLIENT_RESOURCE_EXHAUSTED: + [[fallthrough]]; + case NYdb::EStatus::CLIENT_DEADLINE_EXCEEDED: + [[fallthrough]]; + case NYdb::EStatus::CLIENT_INTERNAL_ERROR: + [[fallthrough]]; + case NYdb::EStatus::CLIENT_CANCELLED: + [[fallthrough]]; + case NYdb::EStatus::CLIENT_UNAUTHENTICATED: + [[fallthrough]]; + case NYdb::EStatus::CLIENT_CALL_UNIMPLEMENTED: + [[fallthrough]]; + case NYdb::EStatus::CLIENT_OUT_OF_RANGE: + [[fallthrough]]; + case NYdb::EStatus::CLIENT_DISCOVERY_FAILED: + [[fallthrough]]; + case NYdb::EStatus::CLIENT_LIMITS_REACHED: + return NErrorCode::ERROR; + } +} + +TYdbSdkCompatibilityConsumer::TYdbSdkCompatibilityConsumer(const TConsumerSettings& settings, + std::shared_ptr<void> destroyEventRef, + TIntrusivePtr<TPQLibPrivate> pqLib, + TIntrusivePtr<ILogger> logger, + NYdb::NPersQueue::TPersQueueClient& client) + : IConsumerImpl(std::move(destroyEventRef), std::move(pqLib)) + , Settings(settings) + , Logger(std::move(logger)) + , Client(client) +{ +} + +TYdbSdkCompatibilityConsumer::~TYdbSdkCompatibilityConsumer() { + Destroy("Destructor called"); +} + +void TYdbSdkCompatibilityConsumer::Init() { + YdbSdkSettings = MakeReadSessionSettings(); // weak_from_this() is used here. +} + +static const THashSet<TString> FederationClusters = { + "iva", + "man", + "myt", + "sas", + "vla", +}; + +static TString FindClusterFromEndpoint(const TString& endpoint) { + size_t pos = endpoint.find(".logbroker.yandex.net"); + if (pos == TString::npos) { + pos = endpoint.find(".logbroker-prestable.yandex.net"); + if (pos == TString::npos) { + return {}; // something strange or cross dc cloud cluster. + } + } + TString prefix = endpoint.substr(0, pos); + if (IsIn(FederationClusters, prefix)) { + return prefix; + } + return {}; // no cluster in cross dc federation. +} + +NYdb::NPersQueue::TReadSessionSettings TYdbSdkCompatibilityConsumer::MakeReadSessionSettings() { + NYdb::NPersQueue::TReadSessionSettings settings; + for (const TString& topic : Settings.Topics) { + settings.AppendTopics(topic); + } + settings.ConsumerName(Settings.ClientId); + settings.MaxMemoryUsageBytes(Settings.MaxMemoryUsage); + if (Settings.MaxTimeLagMs) { + settings.MaxTimeLag(TDuration::MilliSeconds(Settings.MaxTimeLagMs)); + } + if (Settings.ReadTimestampMs) { + settings.StartingMessageTimestamp(TInstant::MilliSeconds(Settings.ReadTimestampMs)); + } + if (Settings.PartitionGroups.size()) { + Y_ENSURE(settings.Topics_.size() == 1); + for (ui64 group : Settings.PartitionGroups) { + settings.Topics_[0].AppendPartitionGroupIds(group); + } + } + if (Settings.ReconnectOnFailure || Settings.ReadFromAllClusterSources) { // ReadFromAllClusterSources implies ReconnectOnFailure and MaxAttempts == inf. + size_t maxRetries = Settings.MaxAttempts; + if (Settings.ReadFromAllClusterSources) { + // Compatibility. + maxRetries = std::numeric_limits<size_t>::max(); + } + + if (Settings.UseV2RetryPolicyInCompatMode) { + settings.RetryPolicy( + NYdb::NPersQueue::IRetryPolicy::GetExponentialBackoffPolicy( + Settings.ReconnectionDelay, + Settings.ReconnectionDelay, + Settings.MaxReconnectionDelay, + maxRetries, + TDuration::Max(), + 2.0, + NYdb::NPersQueue::GetRetryErrorClassV2 + )); + } else { + settings.RetryPolicy( + NYdb::NPersQueue::IRetryPolicy::GetExponentialBackoffPolicy( + Settings.ReconnectionDelay, + Settings.ReconnectionDelay, + Settings.MaxReconnectionDelay, + maxRetries + )); + } + } else { + settings.RetryPolicy(NYdb::NPersQueue::IRetryPolicy::GetNoRetryPolicy()); + } + if (Settings.ReadFromAllClusterSources) { + settings.ReadAll(); + } else { + if (TString cluster = FindClusterFromEndpoint(Settings.Server.Address)) { + if (Settings.ReadMirroredPartitions) { + settings.ReadMirrored(cluster); + } else { + settings.ReadOriginal({ cluster }); + } + } else { + settings.ReadAll(); + } + } + if (Settings.DisableCDS) { + settings.DisableClusterDiscovery(true); + } + + settings.EventHandlers_.HandlersExecutor(NYdb::NPersQueue::CreateThreadPoolExecutorAdapter(PQLib->GetQueuePool().GetQueuePtr(this))); + { + auto weakThis = weak_from_this(); + settings.EventHandlers_.SessionClosedHandler( + [weakThis](const NYdb::NPersQueue::TSessionClosedEvent& event) { + if (auto sharedThis = weakThis.lock()) { + const TString description = event.GetIssues().ToString(); + sharedThis->Destroy(description, ToErrorCode(event.GetStatus())); + } + } + ); + } + + return settings; +} + +NThreading::TFuture<TConsumerCreateResponse> TYdbSdkCompatibilityConsumer::Start(TInstant) noexcept { + ReadSession = Client.CreateReadSession(YdbSdkSettings); + SessionId = ReadSession->GetSessionId(); + DEBUG_LOG("Create read session", "", SessionId); + SubscribeToNextEvent(); + TReadResponse resp; + resp.MutableInit(); + return NThreading::MakeFuture<TConsumerCreateResponse>(TConsumerCreateResponse(std::move(resp))); +} + +NThreading::TFuture<TError> TYdbSdkCompatibilityConsumer::IsDead() noexcept { + return DeadPromise.GetFuture(); +} + +void TYdbSdkCompatibilityConsumer::Destroy(const TError& description) { + if (DeadPromise.HasValue()) { + return; + } + + WARN_LOG("Destroying consumer: " << description, "", SessionId); + + while (!Requests.empty()) { + NPersQueue::TReadResponse resp; + *resp.MutableError() = description; + Requests.front().SetValue(TConsumerMessage(std::move(resp))); + Requests.pop(); + } + + if (ReadSession) { + ReadSession->Close(TDuration::Zero()); + ReadSession = nullptr; + } + + DeadPromise.SetValue(description); + + DestroyPQLibRef(); +} + +void TYdbSdkCompatibilityConsumer::Destroy(const TString& description, NErrorCode::EErrorCode code) { + Destroy(MakeError(description, code)); +} + +void TYdbSdkCompatibilityConsumer::Cancel() { + Destroy(GetCancelReason()); +} + +void TYdbSdkCompatibilityConsumer::GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept { + if (DeadPromise.HasValue()) { + NPersQueue::TReadResponse resp; + *resp.MutableError() = DeadPromise.GetFuture().GetValue(); + promise.SetValue(TConsumerMessage(std::move(resp))); + return; + } + Requests.push(promise); + AnswerToRequests(); + SubscribeToNextEvent(); +} + +void TYdbSdkCompatibilityConsumer::Commit(const TVector<ui64>& cookies) noexcept { + NPersQueue::TReadResponse resp; + for (ui64 cookie : cookies) { + if (cookie >= NextCookie) { + Destroy(TStringBuilder() << "Wrong cookie " << cookie, NErrorCode::WRONG_COOKIE); + return; + } + auto offsetsIt = CookieToOffsets.find(cookie); + if (offsetsIt != CookieToOffsets.end()) { + offsetsIt->second.first.Commit(); + CookiesRequestedToCommit[offsetsIt->second.second].emplace(cookie); + CookieToOffsets.erase(offsetsIt); + } else { + resp.MutableCommit()->AddCookie(cookie); + } + } + if (resp.HasCommit()) { + AddResponse(std::move(resp)); + } +} + +void TYdbSdkCompatibilityConsumer::RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept { + Y_UNUSED(topic); + Y_UNUSED(partition); + auto partitionStreamIt = CurrentPartitionStreams.find(generation); + if (partitionStreamIt != CurrentPartitionStreams.end()) { + partitionStreamIt->second->RequestStatus(); + } +} + +void TYdbSdkCompatibilityConsumer::AnswerToRequests() { + TVector<NYdb::NPersQueue::TReadSessionEvent::TEvent> events; + do { + events = ReadSession->GetEvents(); + for (NYdb::NPersQueue::TReadSessionEvent::TEvent& event : events) { + if (ReadSession) { + std::visit([this](auto&& ev) { return HandleEvent(std::move(ev)); }, event); + } + } + } while (!events.empty() && ReadSession); + + while (!Requests.empty() && !Responses.empty()) { + Requests.front().SetValue(std::move(Responses.front())); + Requests.pop(); + Responses.pop(); + } +} + +void TYdbSdkCompatibilityConsumer::SubscribeToNextEvent() { + if (!SubscribedToNextEvent && !Requests.empty() && ReadSession) { + SubscribedToNextEvent = true; + auto weakThis = weak_from_this(); + auto future = ReadSession->WaitEvent(); + Cerr << "SUBSCRIBING " << ReadSession->GetSessionId() << " future" << future.StateId()->Value() << "\n"; + PQLib->Subscribe(future, this, [weakThis](const NThreading::TFuture<void>&) { + if (auto sharedThis = weakThis.lock()) { + sharedThis->OnReadSessionEvent(); + } + }); + } +} + +void TYdbSdkCompatibilityConsumer::OnReadSessionEvent() { + if (DeadPromise.HasValue()) { + return; + } + + SubscribedToNextEvent = false; + AnswerToRequests(); + SubscribeToNextEvent(); +} + +void TYdbSdkCompatibilityConsumer::HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TDataReceivedEvent&& event) { + if (Settings.UseLockSession) { + auto& offsetSet = PartitionStreamToUncommittedOffsets[event.GetPartitionStream()->GetPartitionStreamId()]; + // Messages could contain holes in offset, but later commit ack will tell us right border. + // So we can easily insert the whole interval with holes included. + // It will be removed from set by specifying proper right border. + offsetSet.InsertInterval(event.GetMessages().front().GetOffset(), event.GetMessages().back().GetOffset() + 1); + } + const ui64 cookie = NextCookie++; + const auto& partitionStream = event.GetPartitionStream(); + auto& offsets = CookieToOffsets[cookie]; + offsets.first.Add(event); + offsets.second = partitionStream->GetPartitionStreamId(); + PartitionStreamToCookies[partitionStream->GetPartitionStreamId()].emplace(cookie); + NPersQueue::TReadResponse resp; + auto& dataResp = *resp.MutableData(); + dataResp.SetCookie(cookie); + auto& batchResp = *dataResp.AddMessageBatch(); + batchResp.SetTopic(MakeLegacyTopicName(partitionStream)); + batchResp.SetPartition(partitionStream->GetPartitionId()); + ui64 maxOffset = 0; + for (auto&& msg : event.GetMessages()) { + auto& msgResp = *batchResp.AddMessage(); + msgResp.SetOffset(msg.GetOffset()); + maxOffset = Max(maxOffset, msg.GetOffset()); + msgResp.SetData(msg.GetData()); + auto& metaResp = *msgResp.MutableMeta(); + metaResp.SetSourceId(msg.GetMessageGroupId()); + metaResp.SetSeqNo(msg.GetSeqNo()); + metaResp.SetCreateTimeMs(msg.GetCreateTime().MilliSeconds()); + metaResp.SetWriteTimeMs(msg.GetWriteTime().MilliSeconds()); + metaResp.SetIp(msg.GetIp()); + for (auto&& [k, v] : msg.GetMeta()->Fields) { + auto& kvResp = *metaResp.MutableExtraFields()->AddItems(); + kvResp.SetKey(k); + kvResp.SetValue(v); + } + } + MaxOffsetToCookie[std::make_pair(partitionStream->GetPartitionStreamId(), maxOffset)] = cookie; + AddResponse(std::move(resp)); +} +void TYdbSdkCompatibilityConsumer::HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TCommitAcknowledgementEvent&& event) { + const ui64 partitionStreamId = event.GetPartitionStream()->GetPartitionStreamId(); + if (Settings.UseLockSession) { + auto& offsetSet = PartitionStreamToUncommittedOffsets[partitionStreamId]; + if (offsetSet.EraseInterval(0, event.GetCommittedOffset())) { // Remove some offsets. + if (offsetSet.Empty()) { // No offsets left. + auto unconfirmedDestroyIt = UnconfirmedDestroys.find(partitionStreamId); + if (unconfirmedDestroyIt != UnconfirmedDestroys.end()) { + // Confirm and forget about this partition stream. + unconfirmedDestroyIt->second.Confirm(); + UnconfirmedDestroys.erase(unconfirmedDestroyIt); + PartitionStreamToUncommittedOffsets.erase(partitionStreamId); + } + } + } + } + const auto offsetPair = std::make_pair(partitionStreamId, event.GetCommittedOffset() - 1); + auto cookieIt = MaxOffsetToCookie.lower_bound(offsetPair); + std::vector<ui64> cookies; + auto end = cookieIt; + auto begin = end; + if (cookieIt != MaxOffsetToCookie.end() && cookieIt->first == offsetPair) { + cookies.push_back(cookieIt->second); + ++end; + } + while (cookieIt != MaxOffsetToCookie.begin()) { + --cookieIt; + if (cookieIt->first.first == partitionStreamId) { + cookies.push_back(cookieIt->second); + begin = cookieIt; + } else { + break; + } + } + if (begin != end) { + MaxOffsetToCookie.erase(begin, end); + } + + NPersQueue::TReadResponse resp; + auto& respCommit = *resp.MutableCommit(); + auto& partitionStreamCookies = PartitionStreamToCookies[partitionStreamId]; + auto& committedCookies = CookiesRequestedToCommit[partitionStreamId]; + for (auto committedCookieIt = cookies.rbegin(), committedCookieEnd = cookies.rend(); committedCookieIt != committedCookieEnd; ++committedCookieIt) { + const ui64 cookie = *committedCookieIt; + respCommit.AddCookie(cookie); + partitionStreamCookies.erase(cookie); + committedCookies.erase(cookie); + } + AddResponse(std::move(resp)); +} + +void TYdbSdkCompatibilityConsumer::HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TCreatePartitionStreamEvent&& event) { + if (Settings.UseLockSession) { + Y_VERIFY(PartitionStreamToUncommittedOffsets[event.GetPartitionStream()->GetPartitionStreamId()].Empty()); + + NPersQueue::TReadResponse resp; + auto& lockResp = *resp.MutableLock(); + const auto& partitionStream = event.GetPartitionStream(); + lockResp.SetTopic(MakeLegacyTopicName(partitionStream)); + lockResp.SetPartition(partitionStream->GetPartitionId()); + lockResp.SetReadOffset(event.GetCommittedOffset()); + lockResp.SetEndOffset(event.GetEndOffset()); + lockResp.SetGeneration(partitionStream->GetPartitionStreamId()); + NThreading::TPromise<TLockInfo> confirmPromise = NThreading::NewPromise<TLockInfo>(); + confirmPromise.GetFuture().Subscribe([event = std::move(event)](const NThreading::TFuture<TLockInfo>& infoFuture) mutable { + const TLockInfo& info = infoFuture.GetValue(); + event.Confirm(info.ReadOffset, info.CommitOffset); + }); // It doesn't matter in what thread this callback will be called. + AddResponse(std::move(resp), std::move(confirmPromise)); + } else { + event.Confirm(); + } +} + +void TYdbSdkCompatibilityConsumer::HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TDestroyPartitionStreamEvent&& event) { + const ui64 partitionStreamId = event.GetPartitionStream()->GetPartitionStreamId(); + CurrentPartitionStreams[partitionStreamId] = event.GetPartitionStream(); + if (Settings.UseLockSession) { + const auto partitionStream = event.GetPartitionStream(); + + Y_VERIFY(UnconfirmedDestroys.find(partitionStreamId) == UnconfirmedDestroys.end()); + if (PartitionStreamToUncommittedOffsets[partitionStreamId].Empty() || Settings.BalanceRightNow) { + PartitionStreamToUncommittedOffsets.erase(partitionStreamId); + event.Confirm(); + } else { + UnconfirmedDestroys.emplace(partitionStreamId, std::move(event)); + } + + NPersQueue::TReadResponse resp; + auto& releaseResp = *resp.MutableRelease(); + releaseResp.SetTopic(MakeLegacyTopicName(partitionStream)); + releaseResp.SetPartition(partitionStream->GetPartitionId()); + releaseResp.SetCanCommit(true); + releaseResp.SetGeneration(partitionStream->GetPartitionStreamId()); + AddResponse(std::move(resp)); + } else { + event.Confirm(); + } +} + +void TYdbSdkCompatibilityConsumer::HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TPartitionStreamStatusEvent&& event) { + NPersQueue::TReadResponse resp; + auto& statusResp = *resp.MutablePartitionStatus(); + const auto& partitionStream = event.GetPartitionStream(); + statusResp.SetGeneration(partitionStream->GetPartitionStreamId()); + statusResp.SetTopic(MakeLegacyTopicName(partitionStream)); + statusResp.SetPartition(partitionStream->GetPartitionId()); + statusResp.SetCommittedOffset(event.GetCommittedOffset()); + statusResp.SetEndOffset(event.GetEndOffset()); + statusResp.SetWriteWatermarkMs(event.GetWriteWatermark().MilliSeconds()); + AddResponse(std::move(resp)); +} + +void TYdbSdkCompatibilityConsumer::HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TPartitionStreamClosedEvent&& event) { + const ui64 partitionStreamId = event.GetPartitionStream()->GetPartitionStreamId(); + CurrentPartitionStreams.erase(partitionStreamId); + + { + NPersQueue::TReadResponse resp; + auto& respCommit = *resp.MutableCommit(); + for (ui64 cookie : CookiesRequestedToCommit[partitionStreamId]) { + respCommit.AddCookie(cookie); + } + if (resp.GetCommit().CookieSize()) { + AddResponse(std::move(resp)); + } + CookiesRequestedToCommit.erase(partitionStreamId); + } + + { + for (ui64 cookie : PartitionStreamToCookies[partitionStreamId]) { + CookieToOffsets.erase(cookie); + } + PartitionStreamToCookies.erase(partitionStreamId); + } + + { + auto begin = MaxOffsetToCookie.lower_bound(std::make_pair(partitionStreamId, 0)); // The first available offset. + if (begin != MaxOffsetToCookie.end() && begin->first.first == partitionStreamId) { + auto end = begin; + while (end != MaxOffsetToCookie.end() && end->first.first == partitionStreamId) { + ++end; + } + MaxOffsetToCookie.erase(begin, end); + } + } + + if (Settings.UseLockSession) { + PartitionStreamToUncommittedOffsets.erase(partitionStreamId); + UnconfirmedDestroys.erase(partitionStreamId); + + if (event.GetReason() != NYdb::NPersQueue::TReadSessionEvent::TPartitionStreamClosedEvent::EReason::DestroyConfirmedByUser) { + NPersQueue::TReadResponse resp; + auto& releaseResp = *resp.MutableRelease(); + const auto& partitionStream = event.GetPartitionStream(); + releaseResp.SetTopic(MakeLegacyTopicName(partitionStream)); + releaseResp.SetPartition(partitionStream->GetPartitionId()); + releaseResp.SetCanCommit(false); + releaseResp.SetGeneration(partitionStream->GetPartitionStreamId()); + AddResponse(std::move(resp)); + } + } +} + +void TYdbSdkCompatibilityConsumer::HandleEvent(NYdb::NPersQueue::TSessionClosedEvent&& event) { + Destroy(event.GetIssues().ToString(), ToErrorCode(event.GetStatus())); +} + +void TYdbSdkCompatibilityConsumer::AddResponse(TReadResponse&& response, NThreading::TPromise<TLockInfo>&& readyToRead) { + Responses.emplace(std::move(response), std::move(readyToRead)); +} + +void TYdbSdkCompatibilityConsumer::AddResponse(TReadResponse&& response) { + Responses.emplace(std::move(response)); +} + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/ydb_sdk_consumer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/ydb_sdk_consumer.h new file mode 100644 index 0000000000..0ff6756686 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/impl/ydb_sdk_consumer.h @@ -0,0 +1,75 @@ +#pragma once +#include "iconsumer_p.h" + +#include <kikimr/public/sdk/cpp/client/ydb_persqueue/persqueue.h> + +#include <library/cpp/containers/disjoint_interval_tree/disjoint_interval_tree.h> + +#include <queue> + +namespace NPersQueue { + +class TYdbSdkCompatibilityConsumer : public IConsumerImpl, + public std::enable_shared_from_this<TYdbSdkCompatibilityConsumer> +{ +public: + TYdbSdkCompatibilityConsumer(const TConsumerSettings& settings, + std::shared_ptr<void> destroyEventRef, + TIntrusivePtr<TPQLibPrivate> pqLib, + TIntrusivePtr<ILogger> logger, + NYdb::NPersQueue::TPersQueueClient& client); + ~TYdbSdkCompatibilityConsumer(); + + void Init() override; + + NThreading::TFuture<TConsumerCreateResponse> Start(TInstant) noexcept override; + NThreading::TFuture<TError> IsDead() noexcept override; + void GetNextMessage(NThreading::TPromise<TConsumerMessage>& promise) noexcept override; + void Commit(const TVector<ui64>& cookies) noexcept override; + void Cancel() override; + void RequestPartitionStatus(const TString& topic, ui64 partition, ui64 generation) noexcept override; + +private: + class TReadSessionEventVisitor; + + NYdb::NPersQueue::TReadSessionSettings MakeReadSessionSettings(); + void SubscribeToNextEvent(); + void OnReadSessionEvent(); + void AnswerToRequests(); + void Destroy(const TError& description); + void Destroy(const TString& description, NErrorCode::EErrorCode code = NErrorCode::ERROR); + + void HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TDataReceivedEvent&& event); + void HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TCommitAcknowledgementEvent&& event); + void HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TCreatePartitionStreamEvent&& event); + void HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TDestroyPartitionStreamEvent&& event); + void HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TPartitionStreamStatusEvent&& event); + void HandleEvent(NYdb::NPersQueue::TReadSessionEvent::TPartitionStreamClosedEvent&& event); + void HandleEvent(NYdb::NPersQueue::TSessionClosedEvent&& event); + + void AddResponse(TReadResponse&& response, NThreading::TPromise<TLockInfo>&& readyToRead); + void AddResponse(TReadResponse&& response); + +private: + const TConsumerSettings Settings; + NYdb::NPersQueue::TReadSessionSettings YdbSdkSettings; + TIntrusivePtr<ILogger> Logger; + TString SessionId; + NYdb::NPersQueue::TPersQueueClient& Client; + std::shared_ptr<NYdb::NPersQueue::IReadSession> ReadSession; + NThreading::TPromise<TError> DeadPromise = NThreading::NewPromise<TError>(); + std::queue<NThreading::TPromise<TConsumerMessage>> Requests; + std::queue<TConsumerMessage> Responses; + bool SubscribedToNextEvent = false; + THashMap<ui64, std::pair<NYdb::NPersQueue::TDeferredCommit, ui64>> CookieToOffsets; // Cookie -> { commit, partition stream id }. + TMap<std::pair<ui64, ui64>, ui64> MaxOffsetToCookie; // { partition stream id, max offset of cookie } -> cookie. + THashMap<ui64, THashSet<ui64>> CookiesRequestedToCommit; // Partition stream id -> cookies that user requested to commit. + THashMap<ui64, THashSet<ui64>> PartitionStreamToCookies; // Partition stream id -> cookies. + ui64 NextCookie = 1; + // Commits for graceful release partition after commit. + THashMap<ui64, TDisjointIntervalTree<ui64>> PartitionStreamToUncommittedOffsets; // Partition stream id -> set of offsets. + THashMap<ui64, NYdb::NPersQueue::TReadSessionEvent::TDestroyPartitionStreamEvent> UnconfirmedDestroys; // Partition stream id -> destroy events. + THashMap<ui64, NYdb::NPersQueue::TPartitionStream::TPtr> CurrentPartitionStreams; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/iprocessor.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/iprocessor.h new file mode 100644 index 0000000000..8b429053fd --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/iprocessor.h @@ -0,0 +1,61 @@ +#pragma once + +#include "types.h" +#include "responses.h" + +#include <library/cpp/threading/future/future.h> + +#include <util/generic/map.h> +#include <util/generic/vector.h> + +namespace NPersQueue { + +struct TPartition { + TString Topic; + ui32 PartitionId = 0; + + bool operator <(const TPartition& other) const { + return std::tie(Topic, PartitionId) < std::tie(other.Topic, other.PartitionId); + } + + explicit TPartition(const TString& topic, const ui32 partitionId) + : Topic(topic) + , PartitionId(partitionId) + {} +}; + +struct TProcessedMessage { + TString Topic; + TString Data; + TString SourceIdPrefix; + ui32 Group = 0; +}; + +struct TProcessedData { + TVector<TProcessedMessage> Messages; // maybe grouped by topic? +}; + +struct TOriginMessage { + TReadResponse::TData::TMessage Message; + NThreading::TPromise<TProcessedData> Processed; +}; + +struct TOriginData { + TMap<TPartition, TVector<TOriginMessage>> Messages; +}; + + +class IProcessor { +public: + virtual ~IProcessor() = default; + + // Returns data and promise for each message. + // Client MUST process each message individually in a deterministic way. + // Each original message CAN be transformed into several messages and written to several topics/sourceids. + // But for any given original message all pairs (topic, sourceIdPrefix) of resulting messages MUST be unique. + // Client MUST finish processing by signaling corresponding promise. + // Otherwise, exactly-once guarantees cannot be hold. + virtual NThreading::TFuture<TOriginData> GetNextData() noexcept = 0; +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/iproducer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/iproducer.h new file mode 100644 index 0000000000..cbb221e9ee --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/iproducer.h @@ -0,0 +1,35 @@ +#pragma once + +#include "types.h" +#include "responses.h" +#include <library/cpp/threading/future/future.h> + +namespace NPersQueue { + +// IProducer is threadsafe. +// If one creates several producers with same SourceId in same topic - first Producer (from point of server view) will die. +// There could be only one producer with concrete SourceId in concrete topic at once. + +class IProducer { +public: + // Start producer. + // Producer can be used after its start will be finished. + virtual NThreading::TFuture<TProducerCreateResponse> Start(TInstant deadline = TInstant::Max()) noexcept = 0; + + NThreading::TFuture<TProducerCreateResponse> Start(TDuration timeout) noexcept { + return Start(TInstant::Now() + timeout); + } + + // Add write request to queue. + virtual NThreading::TFuture<TProducerCommitResponse> Write(TProducerSeqNo seqNo, TData data) noexcept = 0; + + // Add write request to queue without specifying seqNo. So, without deduplication (at least once guarantee). + virtual NThreading::TFuture<TProducerCommitResponse> Write(TData data) noexcept = 0; + + // Get future that is signalled when producer is dead. + virtual NThreading::TFuture<TError> IsDead() noexcept = 0; + + virtual ~IProducer() = default; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/logger.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/logger.h new file mode 100644 index 0000000000..bb3320d795 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/logger.h @@ -0,0 +1,42 @@ +#pragma once +#include <util/generic/ptr.h> +#include <util/generic/string.h> + +namespace NPersQueue { + +class ILogger : public TAtomicRefCount<ILogger> { +public: + virtual ~ILogger() = default; + //level = syslog level + virtual void Log(const TString& msg, const TString& sourceId, const TString& sessionId, int level) = 0; + virtual bool IsEnabled(int level) const = 0; +}; + +class TCerrLogger : public ILogger { +public: + explicit TCerrLogger(int level) + : Level(level) + { + } + + ~TCerrLogger() override = default; + + void Log(const TString& msg, const TString& sourceId, const TString& sessionId, int level) override; + + bool IsEnabled(int level) const override + { + return level <= Level; + } + + int GetLevel() const { + return Level; + } + +private: + static TStringBuf LevelToString(int level); + +private: + int Level; +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h new file mode 100644 index 0000000000..1b1f7646c8 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h @@ -0,0 +1,50 @@ +#pragma once + +#include "types.h" +#include "iproducer.h" +#include "iconsumer.h" +#include "iprocessor.h" +#include "responses.h" + +#include <library/cpp/threading/future/future.h> + +#include <util/generic/noncopyable.h> +#include <util/generic/ptr.h> + +namespace NPersQueue { + +class TPQLibPrivate; + +class TPQLib : public TNonCopyable { +public: + explicit TPQLib(const TPQLibSettings& settings = TPQLibSettings()); + ~TPQLib(); + + // Deprecated flag means that PQLib doen't wait its objects during destruction. + // This behaviour leads to potential thread leakage and thread sanitizer errors. + // It is recommended to specify deprecated = false. In that case PQLib will cancel + // all objects and correctly wait for its threads in destructor. But, unfortunately, + // we can't migrate all current clients to this default behaviour automatically, + // because this can break existing programs unpredictably. + + // Producers creation + THolder<IProducer> CreateProducer(const TProducerSettings& settings, TIntrusivePtr<ILogger> logger = nullptr, bool deprecated = false); + THolder<IProducer> CreateMultiClusterProducer(const TMultiClusterProducerSettings& settings, TIntrusivePtr<ILogger> logger = nullptr, bool deprecated = false); + + // Consumers creation + THolder<IConsumer> CreateConsumer(const TConsumerSettings& settings, TIntrusivePtr<ILogger> logger = nullptr, bool deprecated = false); + + // Processors creation + THolder<IProcessor> CreateProcessor(const TProcessorSettings& settings, TIntrusivePtr<ILogger> logger = nullptr, bool deprecated = false); + + void SetLogger(TIntrusivePtr<ILogger> logger); + + TString GetUserAgent() const; + void SetUserAgent(const TString& userAgent); + +private: + TIntrusivePtr<TPQLibPrivate> Impl; + TAtomic Alive = 1; // Debug check for valid PQLib usage. +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h new file mode 100644 index 0000000000..de5a801172 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/responses.h @@ -0,0 +1,106 @@ +#pragma once + +#include "types.h" + +#include <library/cpp/threading/future/future.h> +#include <util/generic/string.h> + +namespace NPersQueue { + +struct TProducerCreateResponse { + TProducerCreateResponse(TWriteResponse&& response) + : Response(std::move(response)) + { + } + + TWriteResponse Response; +}; + +struct TProducerCommitResponse { + TProducerCommitResponse(TProducerSeqNo seqNo, TData data, TWriteResponse&& response) + : Response(std::move(response)) + , Data(std::move(data)) + , SeqNo(seqNo) + { + } + + TWriteResponse Response; + TData Data; + TProducerSeqNo SeqNo; +}; + + +struct TConsumerCreateResponse { + TConsumerCreateResponse(TReadResponse&& response) + : Response(std::move(response)) + { + } + + //will contain Error or Init + TReadResponse Response; +}; + +enum EMessageType { + EMT_LOCK, + EMT_RELEASE, + EMT_DATA, + EMT_ERROR, + EMT_STATUS, + EMT_COMMIT +}; + +struct TLockInfo { + ui64 ReadOffset = 0; + ui64 CommitOffset = 0; + bool VerifyReadOffset = false; + + TLockInfo() = default; + + // compatibility with msvc2015 + TLockInfo(ui64 readOffset, ui64 commitOffset, bool verifyReadOffset) + : ReadOffset(readOffset) + , CommitOffset(commitOffset) + , VerifyReadOffset(verifyReadOffset) + {} +}; + +static EMessageType GetType(const TReadResponse& response) { + if (response.HasData()) return EMT_DATA; + if (response.HasError()) return EMT_ERROR; + if (response.HasCommit()) return EMT_COMMIT; + if (response.HasPartitionStatus()) return EMT_STATUS; + if (response.HasRelease()) return EMT_RELEASE; + if (response.HasLock()) return EMT_LOCK; + + // for no warn return anything. + return EMT_LOCK; +} + + + +struct TConsumerMessage { + EMessageType Type; + //for DATA/ERROR/LOCK/RELEASE/COMMIT/STATUS: + //will contain Error, Data, Lock or Release and be consistent with Type + TReadResponse Response; + + //for LOCK only: + mutable NThreading::TPromise<TLockInfo> ReadyToRead; + + TConsumerMessage(TReadResponse&& response, NThreading::TPromise<TLockInfo>&& readyToRead) + : Type(EMT_LOCK) + , Response(std::move(response)) + , ReadyToRead(std::move(readyToRead)) + { + Y_VERIFY(GetType(Response) == EMT_LOCK); + } + + + explicit TConsumerMessage(TReadResponse&& response) + : Type(GetType(response)) + , Response(std::move(response)) + { + } +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/samples/consumer/main.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/samples/consumer/main.cpp new file mode 100644 index 0000000000..b30115e294 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/samples/consumer/main.cpp @@ -0,0 +1,236 @@ +#include <util/datetime/base.h> +#include <util/stream/file.h> +#include <util/string/join.h> +#include <util/string/vector.h> + +#include <google/protobuf/message.h> + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> + +using namespace NPersQueue; + + +int main() { + + // [BEGIN create pqlib] + // Step 1. First create logger. It could be your logger inherited from ILogger interface or one from lib. + TIntrusivePtr<ILogger> logger = new TCerrLogger(7); + + // Step 2. Create settings of main TPQLib object. + TPQLibSettings pqLibSettings; + + pqLibSettings.DefaultLogger = logger; + + // Step 3. Create TPQLib object. + TPQLib pq(pqLibSettings); + // [END create pqlib] + + // [BEGIN create consumer] + // Step 4. Parameters used in sample. + + // List of topics to read. + TVector<TString> topics = {"topic_path"}; + + // Path of consumer that will read topics. + TString consumerPath = "consumer_path"; + + // OAuth token. + TString oauthToken = "oauth_token"; + + // IAM JWT key + TString jwtKey = "<jwt json key>"; + + // TVM parameters. + TString tvmSecret = "tvm_secret"; + ui32 tvmClientId = 200000; + ui32 tvmDstId = 2001147; // logbroker main prestable + tvmDstId = 2001059; // logbroker main production + + // Logbroker endpoint. + TString endpoint = "man.logbroker.yandex.net"; + + // Simple single read. + + // Step 5. Create settings of consumer object. + TConsumerSettings settings; + // Fill consumer settings. See types.h for more information. + // Settings without defaults, must be set always. + settings.Server.Address = endpoint; + { + // Working with logbroker in Yandex.Cloud: + // Fill database for Yandex.Cloud Logbroker. + settings.Server.Database = "/Root"; + // Enable TLS. + settings.Server.EnableSecureConnection(""); + } + settings.Topics = topics; + settings.ClientId = consumerPath; + + // One must set CredentialsProvider for authentification. There could be OAuth and TVM providers. + // Create OAuth provider from OAuth ticket. + settings.CredentialsProvider = CreateOAuthCredentialsProvider(oauthToken); + // Or create TVM provider from TVM settings. + settings.CredentialsProvider = CreateTVMCredentialsProvider(tvmSecret, tvmClientId, tvmDstId, logger); + // Or create IAM provider from jwt-key + settings.CredentialsProvider = CreateIAMJwtParamsCredentialsForwarder(jwtKey, logger); + + // Step 6. Create consumer object. PQLib object must be valid during all lifetime of consumer, otherwise reading from consumer will fail but this is not UB. + auto consumer = pq.CreateConsumer(settings, logger); + + // Step 7. Start consuming data. + auto future = consumer->Start(); + + // Step 9. Wait initialization and check it result. + future.Wait(); + + if (future.GetValue().Response.HasError()) { + return 1; + } + // [END create consumer] + + // [BEGIN read messages] + // Step 10. Simple single read. + // Step 10.1. Request message from consumer. + auto msg = consumer->GetNextMessage(); + + // Step 10.2. Wait for message. + msg.Wait(); + + auto value = msg.GetValue(); + + // Step 10.3. Check result. + switch (value.Type) { + // Step 10.3.1. Result is error - consumer is in incorrect state now. You can log error and recreate everything starting from step 5. + case EMT_ERROR: + return 1; + // Step 10.3.2. Read data. + case EMT_DATA: { + // Process received messages. Inside will be several batches. + for (const auto& t : value.Response.data().message_batch()) { + const TString topic = t.topic(); + const ui32 partition = t.partition(); + Y_UNUSED(topic); + Y_UNUSED(partition); + // There are several messages from partition of topic inside one batch. + for (const auto& m : t.message()) { + const ui64 offset = m.offset(); + const TString& data = m.data(); + // Each message has offset in partition, payload data (decompressed or as is) and metadata. + Y_UNUSED(offset); + Y_UNUSED(data); + } + } + + // Step 10.3.3. Each processed data batch from step 10.3.2 must be committed once all data inside is processed. + + // This must be done using cookie from data batch. + auto cookie = msg.GetValue().Response.GetData().GetCookie(); + // Commit cookie['s]. Acknowledge from server will be received inside one of next messages from consumer later. + consumer->Commit({cookie}); + break; + } + // Step 10.3.4. Commit acknowledge. Can be used for logging and debugging purposes. (Can't be first message, only possible if one makes several + // consecutive reads. + case EMT_COMMIT: + break; + // Handle other cases - only for compiling, there could not be any. + default: + break; + + } + // [END read messages] + + // [BEGIN advanced read messages] + + // Point that consumer will perform advanced reads. + settings.UseLockSession = true; + + { + // Step 11. Create consumer object. PQLib object must be valid during all lifetime of consumer, otherwise reading from consumer will fail but this is not UB. + auto consumer = pq.CreateConsumer(settings, logger); + + // Step 12. Start consuming data. + auto future = consumer->Start(); + + // Step 14. Wait initialization and check it result. + future.Wait(); + + if (future.GetValue().Response.HasError()) { + return 1; + } + + // Step 14. Event loop. + while (true) { + // Step 14.1. Request message from consumer. + auto msg = consumer->GetNextMessage(); + + // Step 14.2. Wait for message. + msg.Wait(); + + const auto& value = msg.GetValue(); + + // Step 14.3. Check result. + switch (value.Type) { + // Step 14.3.1. Result is error: consumer is in incorrect state now. You can log error and recreate everything starting from step 11. + case EMT_ERROR: + return 1; + // Step 14.3.2. Lock request from server. + case EMT_LOCK: { + const TString& topic = value.Response.lock().topic(); + const ui32 partition = value.Response.lock().partition(); + // Server is ready to send data from corresponding partition of topic. Topic is legacy topic name = rt3.<cluster>--<converted topic path> + Y_UNUSED(topic); + Y_UNUSED(partition); + // One must react on this event with setting promise ReadyToRead right now or later. Only after this client will start receiving data from this partition. + msg.GetValue().ReadyToRead.SetValue(TLockInfo{0 /*readOffset*/, 0 /*commitOffset*/, false /*verifyReadOffset*/}); + break; + } + // Step 14.3.3. Release request from server. + case EMT_RELEASE: { + const TString& topic = value.Response.release().topic(); + const ui32 partition = value.Response.release().partition(); + Y_UNUSED(topic); + Y_UNUSED(partition); + // // Server will not send any data from partition of topic until next Lock request is confirmed by client. + break; + } + // Step 14.3.4. Read data. + case EMT_DATA: { + // Process received messages. Inside will be several batches. The same as in the simple read. + for (const auto& t : value.Response.data().message_batch()) { + const TString topic = t.topic(); + const ui32 partition = t.partition(); + Y_UNUSED(topic); + Y_UNUSED(partition); + // There are several messages from partition of topic inside one batch. + for (const auto& m : t.message()) { + const ui64 offset = m.offset(); + const TString& data = m.data(); + // Each message has offset in partition, payload data (decompressed or as is) and metadata. + Y_UNUSED(offset); + Y_UNUSED(data); + } + } + + // Step 14.3.5. Each processed data batch from step 14.3.4 should be committed after all data inside is processed. + + // This must be done using cookie from data batch. + auto cookie = msg.GetValue().Response.GetData().GetCookie(); + // Commit cookie['s]. Acknowledge from server will be received inside message from consumer later. + consumer->Commit({cookie}); + break; + } + // Step 14.3.6. Commit acknowledge. Can be used for logging and debugging purposes. + case EMT_STATUS: + case EMT_COMMIT: + break; + } // switch() + } // while() + } + + // [END advanced read messages] + + // Step 15. End of working. Nothing special needs to be done - desctuction of objects only. + + return 0; +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/samples/producer/main.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/samples/producer/main.cpp new file mode 100644 index 0000000000..c91838892b --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/samples/producer/main.cpp @@ -0,0 +1,174 @@ +#include <queue> + +#include <util/datetime/base.h> +#include <util/stream/file.h> +#include <util/string/join.h> +#include <util/string/vector.h> + +#include <google/protobuf/message.h> + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> + +using namespace NPersQueue; + + +int main() { + + + // [BEGIN create pqlib] + // Step 1. First create logger. It could be your logger inherited from ILogger interface or one from lib. + TIntrusivePtr<ILogger> logger = new TCerrLogger(7); + + // Step 2. Create settings of main TPQLib object. + TPQLibSettings pqLibSettings; + + // Number of threads for processing gRPC events. + pqLibSettings.ThreadsCount = 1; + // Number of threads for compression/decomression actions. + pqLibSettings.CompressionPoolThreads = 3; + // Default logger. Used inside consumers and producers if no logger provided. + pqLibSettings.DefaultLogger = logger; + + // Step 3. Create TPQLib object. + TPQLib pq(pqLibSettings); + // [END create pqlib] + + // [BEGIN create producer] + // Step 4. Initialize producer parameters + + // Topics to read. + TString topic = "topic_path"; + // Partitions group. 0 means any. + ui32 group = 0; + + // Source identifier of messages group. + TString sourceId = "source_id"; + + // OAuth token. + TString oauthToken = "oauth_token"; + + // IAM JWT key + TString jwtKey = "<jwt json key>"; + + // TVM parameters. + TString tvmSecret = "tvm_secret"; + ui32 tvmClientId = 200000; + ui32 tvmDstId = 2001147; // logbroker main prestable + tvmDstId = 2001059; // logbroker main production + + // Logbroker endpoint. + TString endpoint = "logbroker.yandex.net"; + + // Step 5. Put producer parameters into settings. + TProducerSettings settings; + // Logbroker endpoint. + settings.Server.Address = endpoint; + { + // Working with logbroker in Yandex.Cloud: + // Fill database for Yandex.Cloud Logbroker. + settings.Server.Database = "/Root"; + // Enable TLS. + settings.Server.EnableSecureConnection(""); + } + // Topic path. + settings.Topic = topic; + // SourceId of messages group. + settings.SourceId = sourceId; + // Partition group to write to. + settings.PartitionGroup = group; + // ReconnectOnFailure will gurantee that producer will retry writes in case of failre. + settings.ReconnectOnFailure = true; //retries on server errors + // Codec describes how to compress data. + settings.Codec = NPersQueueCommon::ECodec::GZIP; + + // One must set CredentialsProvider for authentification. There could be OAuth and TVM providers. + // Create OAuth provider from OAuth ticket. + settings.CredentialsProvider = CreateOAuthCredentialsProvider(oauthToken); + // Or create TVM provider from TVM settings. + settings.CredentialsProvider = CreateTVMCredentialsProvider(tvmSecret, tvmClientId, tvmDstId, logger); + // Or create IAM provider from jwt-key + settings.CredentialsProvider = CreateIAMJwtParamsCredentialsForwarder(jwtKey, logger); + + // Step 6. Create producer object. PQLib object must be valid during all lifetime of producer, otherwise writing to producer will fail but this is not UB. + auto producer = pq.CreateProducer(settings, logger); + + // Step 7. Start producer. + auto future = producer->Start(); + + // Step 8. Wait initialization and check it result. + future.Wait(); + + if (future.GetValue().Response.HasError()) { + return 1; + } + // [END create producer] + + // [BEGIN write message] + + // Step 9. Single write sample. + // Step 9.1. Form message. + // Creation timestamp. + TInstant timestamp = TInstant::Now(); + // Message payload. + TString payload = "abacaba"; + + TData data{payload, timestamp}; + // Corresponding messages's sequence number. + ui64 seqNo = 1; + + // Step 9.2. Write message. Until result future is signalled - message is stored inside PQLib memory. Future will be signalled with message payload and result of writing. + auto result = producer->Write(seqNo, data); + + // Step 9.3 Check writing result. + // Check all results. + result.Wait(); + Y_VERIFY(!result.GetValue().Response.HasError()); + + // Only when future is destroyed no memory is used by message. + + // [END write message] + + // [BEGIN write messages with inflight] + + // Step 10. Writing loop with in-flight. + // Max in-flight border, parameter. + ui32 maxInflyCount = 10; + // Queue for messages in-flight. + std::queue<NThreading::TFuture<TProducerCommitResponse>> v; + + while(true) { + // Creation timestamp. + TInstant timestamp = TInstant::Now(); + // Message payload. + TString payload = "abacaba"; + TData data{payload, timestamp}; + + // SeqNo must increase for consecutive messages. + ++seqNo; + + // Step 10.1. Write message. + auto result = producer->Write(seqNo, data); + // Step 10.2. Put future inside queue. + v.push(result); + // Step 10.3. If in-flight size is too big. + if (v.size() > maxInflyCount) { + // Step 10.4. Wait for first message in queue processing result, check it. + v.front().Wait(); + Y_VERIFY(!v.front().GetValue().Response.HasError()); + // Step 10.5. Drop first processed message. Only after memory of this message will be freed. + v.pop(); + } + } + + // Step 10.6. After processing loop one must wait for all not yet processed messages. + + while (!v.empty()) { + v.front().Wait(); + Y_VERIFY(!v.front().GetValue().Response.HasError()); + v.pop(); + } + + // [END write messages with inflight] + + return 0; +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/types.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/types.h new file mode 100644 index 0000000000..37c4423d6f --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/types.h @@ -0,0 +1,402 @@ +#pragma once +#include "credentials_provider.h" +#include "logger.h" +#include <kikimr/yndx/api/protos/persqueue.pb.h> + + +#include <library/cpp/logger/priority.h> + +#include <util/generic/string.h> +#include <util/generic/buffer.h> +#include <util/generic/maybe.h> +#include <util/generic/vector.h> +#include <util/generic/hash.h> +#include <util/datetime/base.h> +#include <util/stream/output.h> +#include <util/string/builder.h> + +#include <grpc++/channel.h> + +#include <limits> + +namespace NPersQueue { +using ECodec = NPersQueueCommon::ECodec; +using TCredentials = NPersQueueCommon::TCredentials; +using TError = NPersQueueCommon::TError; + +using TProducerSeqNo = ui64; // file offset or something with equivalent meaning. +using TConsumerOffset = i64; // partition offset. Have special value -1. All other values are >= 0. +using TCredProviderPtr = std::shared_ptr<ICredentialsProvider>; + +constexpr ui32 DEFAULT_MEMORY_USAGE = 100 << 20; //100mb +constexpr ui32 DEFAULT_READS_INFLY = 10; +constexpr ui32 DEFAULT_BATCH_SIZE = 1 << 20; //1mb +constexpr TDuration DEFAULT_CHANNEL_CREATION_TIMEOUT = TDuration::Seconds(10); +constexpr TDuration DEFAULT_RECONNECTION_DELAY = TDuration::MilliSeconds(10); +constexpr TDuration DEFAULT_MAX_RECONNECTION_DELAY = TDuration::Seconds(1); +constexpr TDuration DEFAULT_START_SESSION_TIMEOUT = TDuration::Seconds(30); +constexpr TDuration DEFAULT_CHECK_TOKEN_PERIOD = TDuration::Hours(1); +inline const THashMap<ECodec, TString> CodecIdByType{{ECodec::DEFAULT, TString(1, '\0')}, {ECodec::RAW, TString(1, '\0')}, {ECodec::GZIP, TString(1, '\1')}, {ECodec::LZOP, TString(1, '\2')}, {ECodec::ZSTD, TString(1, '\3')}}; + +struct TPQLibSettings { + size_t ThreadsCount = 1; + size_t CompressionPoolThreads = 2; + size_t GRpcThreads = 1; + TDuration ChannelCreationTimeout = DEFAULT_CHANNEL_CREATION_TIMEOUT; + TIntrusivePtr<ILogger> DefaultLogger = {}; + + bool EnableGRpcV1 = std::getenv("PERSQUEUE_GRPC_API_V1_ENABLED"); +}; + +enum class EClusterDiscoveryUsageMode { + Auto, + Use, + DontUse, +}; + +// destination server +struct TServerSetting { + // Endpoint of logbroker cluster to connect. + // Possible variants: + // + // 1. lbkx cluster. + // Set address to lbkx.logbroker.yandex.net + // + // 2. logbroker federation. + // If you want to write/read to/from concrete cluster + // (with ReadMirroredPartitions or not, but without ReadFromAllClusterSources option), + // use the following addresses: + // iva.logbroker.yandex.net + // man.logbroker.yandex.net + // myt.logbroker.yandex.net + // sas.logbroker.yandex.net + // vla.logbroker.yandex.net + // + // If you create consumer with ReadFromAllClusterSources option, + // use logbroker.yandex.net + TString Address; + + ui16 Port = 2135; + TString Database = GetRootDatabase(); + EClusterDiscoveryUsageMode UseLogbrokerCDS = EClusterDiscoveryUsageMode::Auto; + bool UseSecureConnection = false; + TString CaCert = TString(); + + TServerSetting(TString address, ui16 port = 2135) + : Address(std::move(address)) + , Port(port) + { + Y_VERIFY(!Address.empty()); + Y_VERIFY(Port > 0); + } + + TServerSetting() = default; + TServerSetting(const TServerSetting&) = default; + TServerSetting(TServerSetting&&) = default; + + TServerSetting& operator=(const TServerSetting&) = default; + TServerSetting& operator=(TServerSetting&&) = default; + + TString GetFullAddressString() const { + return TStringBuilder() << Address << ':' << Port; + } + void EnableSecureConnection(const TString& caCert) { + UseSecureConnection = true; + CaCert = caCert; + } + void DisableSecureConnection() { + UseSecureConnection = false; + CaCert.clear(); + } + + const TString& GetRootDatabase(); + + bool operator < (const TServerSetting& rhs) const { + return std::tie(Address, Port, Database, UseSecureConnection) < std::tie(rhs.Address, rhs.Port, rhs.Database, rhs.UseSecureConnection); + } +}; + +// +// consumer/producer settings +// + +struct TProducerSettings { + TServerSetting Server; + std::shared_ptr<ICredentialsProvider> CredentialsProvider = CreateInsecureCredentialsProvider(); + + TString Topic; + TString SourceId; + + ui32 PartitionGroup = 0; + ECodec Codec = ECodec::GZIP; + int Quality = -1; + THashMap<TString, TString> ExtraAttrs; + + // Should we reconnect when session breaks? + bool ReconnectOnFailure = false; + + // + // Setting for reconnection + // + + // In case of lost connection producer will try + // to reconnect and resend data <= MaxAttempts times. + // Then producer will be dead. + unsigned MaxAttempts = std::numeric_limits<unsigned>::max(); + + // Time delay before next producer reconnection attempt. + // The second attempt will be in 2 * ReconnectionDelay, + // the third - in 4 * ReconnectionDelay, etc. Maximum value + // of delay is MaxReconnectionDelay. + TDuration ReconnectionDelay = DEFAULT_RECONNECTION_DELAY; + TDuration MaxReconnectionDelay = DEFAULT_MAX_RECONNECTION_DELAY; + + // Timeout for session initialization operation. + TDuration StartSessionTimeout = DEFAULT_START_SESSION_TIMEOUT; + // Producer will check token from 'ICredentialsProvider' with this period and if changed - send to server as soon as possible. + // Available only with gRPC data plane API v1. + TDuration CheckTokenPeriod = DEFAULT_CHECK_TOKEN_PERIOD; + + // Preferred cluster to write. + // This setting is used only in case of cluster discovery is used. + // Case insensitive. + // If PreferredCluster is not available, + // producer will connect to other cluster. + TString PreferredCluster; + + bool DisableCDS = false; // only for grpcv1 mode +}; + +struct TMultiClusterProducerSettings { + struct TServerWeight { + // Setting of producer connected to concrete DC. + // ReconnectOnFailure must be true. + // MaxAttempts must be a finite number. + // SourceId will be complemented with source id prefix and producer type. Can be empty. + // There is no guarantee that messages will arrive in the order of writes in case of one DC is absent. + TProducerSettings ProducerSettings; + unsigned Weight; + + explicit TServerWeight(TProducerSettings producerSettings, unsigned weight = 1) + : ProducerSettings(std::move(producerSettings)) + , Weight(weight) + { + Y_VERIFY(Weight > 0); + } + + TServerWeight() = default; + TServerWeight(const TServerWeight&) = default; + TServerWeight(TServerWeight&&) = default; + + TServerWeight& operator=(const TServerWeight&) = default; + TServerWeight& operator=(TServerWeight&&) = default; + }; + + std::vector<TServerWeight> ServerWeights; + size_t MinimumWorkingDcsCount = 1; // Minimum Dcs to start and not to die. + TString SourceIdPrefix; // There is no guarantee that messages will arrive in the order of writes in case of one DC is absent. +}; + +struct TConsumerSettings { + TVector<TString> Topics; + bool ReadMirroredPartitions = true; + + // Discovers all clusters where topic is presented and create consumer that reads from all of them. + // Conflicts with ReadMirroredPartitions option (should be set to "no"). + // Will ignore the following options: + // ReconnectOnFailure (forced to "yes") + // MaxAttempts (forced to infinity == std::numeric_limits<unsigned>::max()) + // Server.UseLogbrokerCDS (forced to "Use") + // + // Also see description of option ReconnectOnFailure. + bool ReadFromAllClusterSources = false; + + TServerSetting Server; + std::shared_ptr<ICredentialsProvider> CredentialsProvider = CreateInsecureCredentialsProvider(); + TString ClientId; + + bool UseLockSession = false; + ui32 PartsAtOnce = 0; + + bool Unpack = true; + bool SkipBrokenChunks = false; + + ui64 MaxMemoryUsage = DEFAULT_MEMORY_USAGE; + ui32 MaxInflyRequests = DEFAULT_READS_INFLY; + + // read settings + ui32 MaxCount = 0; // zero means unlimited + ui32 MaxSize = DEFAULT_BATCH_SIZE; // zero means unlimited + ui32 MaxTimeLagMs = 0; // zero means unlimited + ui64 ReadTimestampMs = 0; // read data from this timestamp + + ui32 MaxUncommittedCount = 0; // zero means unlimited + ui64 MaxUncommittedSize = 0; // zero means unlimited + + TVector<ui32> PartitionGroups; // groups to read + + bool PreferLocalProxy = false; + bool BalanceRightNow = false; + bool CommitsDisabled = false; + + // Should we reconnect when session breaks? + // + // Keep in mind that it is possible to receive + // message duplicates when reconnection occures. + bool ReconnectOnFailure = false; + + // + // Settings for reconnection + // + + // In case of lost connection consumer will try + // to reconnect and resend requests <= MaxAttempts times. + // Then consumer will be dead. + unsigned MaxAttempts = std::numeric_limits<unsigned>::max(); + + // Time delay before next consumer reconnection attempt. + // The second attempt will be in 2 * ReconnectionDelay, + // the third - in 3 * ReconnectionDelay, etc. Maximum value + // of delay is MaxReconnectionDelay. + TDuration ReconnectionDelay = DEFAULT_RECONNECTION_DELAY; + TDuration MaxReconnectionDelay = DEFAULT_MAX_RECONNECTION_DELAY; + TDuration StartSessionTimeout = DEFAULT_START_SESSION_TIMEOUT; + + bool DisableCDS = false; // only for grpcv1 mode + + bool UseV2RetryPolicyInCompatMode = false; +}; + +struct TProcessorSettings { + ui64 MaxOriginDataMemoryUsage; + ui64 MaxProcessedDataMemoryUsage; + TDuration SourceIdIdleTimeout; + + std::shared_ptr<ICredentialsProvider> CredentialsProvider; + + //UseLockSession parameter will be ignored + TConsumerSettings ConsumerSettings; + + //Topic and SourceId pararameters will be ignored + TProducerSettings ProducerSettings; + + TProcessorSettings() + : MaxOriginDataMemoryUsage(DEFAULT_MEMORY_USAGE) + , MaxProcessedDataMemoryUsage(DEFAULT_MEMORY_USAGE) + , SourceIdIdleTimeout(TDuration::Hours(1)) + , CredentialsProvider(CreateInsecureCredentialsProvider()) + { + ConsumerSettings.CredentialsProvider = ProducerSettings.CredentialsProvider = CredentialsProvider; + } +}; + +class TData { +public: + // empty: not for passing to producer + TData() = default; + + // implicit constructor from source data to be encoded later by producer with default codec from settings + // or explicitly specified one + TData(TString sourceData, ECodec codec = ECodec::DEFAULT, TInstant timestamp = TInstant::Now()) + : TData(timestamp, sourceData, codec == ECodec::RAW ? sourceData : TString(), codec, sourceData.size()) + { + Y_VERIFY(!SourceData.empty()); + Y_VERIFY(timestamp != TInstant::Zero()); + } + + TData(TString sourceData, TInstant timestamp) + : TData(std::move(sourceData), ECodec::DEFAULT, timestamp) + { + } + + // construct already encoded TData + static TData Encoded(TString encodedData, ECodec codec, TInstant timestamp = TInstant::Now(), ui32 originalSize = 0) { + Y_VERIFY(!encodedData.empty()); + Y_VERIFY(codec != ECodec::RAW && codec != ECodec::DEFAULT); + Y_VERIFY(timestamp != TInstant::Zero()); + return TData(timestamp, TString(encodedData), encodedData, codec, originalSize == 0 ? encodedData.size() * 10 : originalSize); + } + + // construct user defined raw data that shouldn't be encoded by producer + static TData Raw(TString rawData, TInstant timestamp = TInstant::Now()) { + Y_VERIFY(!rawData.empty()); + Y_VERIFY(timestamp != TInstant::Zero()); + return TData(timestamp, rawData, rawData, ECodec::RAW, rawData.size()); + } + + TData(const TData&) = default; + TData(TData&&) = default; + + TData& operator=(const TData&) = default; + TData& operator=(TData&&) = default; + + ECodec GetCodecType() const { + return Codec; + } + + bool IsEncoded() const { + return !EncodedData.empty(); + } + + const TString& GetEncodedData() const { + Y_VERIFY(IsEncoded()); + return EncodedData; + } + + // gets some data for error report + const TString& GetSourceData() const { + return SourceData; + } + + TInstant GetTimestamp() const { + return Timestamp; + } + + bool Empty() const { + return GetTimestamp() == TInstant::Zero(); + } + + // encoding of data by default codec or data codec that is not encoded + // if data's codec is default, uses defaultCodec parameter, + // otherwise uses data's codec + static TData Encode(TData data, ECodec defaultCodec, int quality = -1); + + ui32 GetOriginalSize() const { + return OriginalSize; + } + + // special function to make raw enconding that doesn't need cpu time. + static TData MakeRawIfNotEncoded(TData data); + + bool operator==(const TData& data) const { + return Codec == data.Codec + && Timestamp == data.Timestamp + && SourceData == data.SourceData + && EncodedData == data.EncodedData; + } + + bool operator!=(const TData& data) const { + return !operator==(data); + } + +private: + TData(TInstant timestamp, TString sourceData, TString encodedData, ECodec codec, ui32 originalSize) + : Timestamp(timestamp) + , SourceData(std::move(sourceData)) + , EncodedData(std::move(encodedData)) + , Codec(codec) + , OriginalSize(originalSize) + { + } + + static THolder<IOutputStream> CreateCoder(ECodec codec, TData& data, int quality); + +private: + TInstant Timestamp; + TString SourceData; + TString EncodedData; + ECodec Codec = ECodec::DEFAULT; + ui32 OriginalSize = 0; +}; + +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/CMakeLists.txt b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/CMakeLists.txt new file mode 100644 index 0000000000..ec1f7fb67f --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/CMakeLists.txt @@ -0,0 +1,30 @@ + +# This file was gererated by the build system used internally in the Yandex monorepo. +# Only simple modifications are allowed (adding source-files to targets, adding simple properties +# like target_include_directories). These modifications will be ported to original +# ya.make files by maintainers. Any complex modifications which can't be ported back to the +# original buildsystem will not be accepted. + + + +add_library(cpp-v2-ut_utils) +target_compile_options(cpp-v2-ut_utils PRIVATE + -DUSE_CURRENT_UDF_ABI_VERSION +) +target_link_libraries(cpp-v2-ut_utils PUBLIC + contrib-libs-cxxsupp + yutil + deprecated-cpp-v2 + yndx-grpc_services-persqueue + yndx-persqueue-msgbus_server + cpp-grpc-server + cpp-testing-unittest + core-client-server + ydb-core-testlib + library-persqueue-topic_parser_public +) +target_sources(cpp-v2-ut_utils PRIVATE + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/data_writer.cpp + ${CMAKE_SOURCE_DIR}/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_pqlib.cpp +) diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/data_writer.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/data_writer.cpp new file mode 100644 index 0000000000..6f201f39f5 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/data_writer.cpp @@ -0,0 +1,263 @@ +#include "data_writer.h" + +namespace NPersQueue::NTests { +using namespace NKikimr; + +void TPQDataWriter::Read( + const TString& topic, const TString& clientId, const TString& ticket, bool error, bool checkACL, + bool useBatching, bool onlyCreate +) { + + grpc::ClientContext context; + + auto stream = StubP_->ReadSession(&context); + UNIT_ASSERT(stream); + + // Send initial request. + TReadRequest req; + TReadResponse resp; + + req.MutableInit()->AddTopics(topic); + + req.MutableInit()->SetClientId(clientId); + req.MutableInit()->SetProxyCookie(ProxyCookie); + if (!ticket.empty()) { + req.MutableCredentials()->SetTvmServiceTicket(ticket); + } + + if (useBatching) { + req.MutableInit()->SetProtocolVersion(TReadRequest::Batching); + } + + if (!stream->Write(req)) { + ythrow yexception() << "write fail"; + } + //TODO[komels]: why this leads to timeout? + //Server.AnnoyingClient->GetClientInfo({topic}, clientId, true); + UNIT_ASSERT(stream->Read(&resp)); + if (error) { + UNIT_ASSERT(resp.HasError()); + return; + } + UNIT_ASSERT_C(resp.HasInit(), resp); + + if (onlyCreate) + return; + + for (ui32 i = 0; i < 11; ++i) { + TReadRequest req; + + req.MutableRead()->SetMaxCount(1); + + if (!stream->Write(req)) { + ythrow yexception() << "write fail"; + } + Server.AnnoyingClient->AlterTopic(FullTopicName, i < 10 ? 2 : 3); + + } + + if (checkACL) { + NACLib::TDiffACL acl; + acl.RemoveAccess(NACLib::EAccessType::Allow, NACLib::SelectRow, clientId + "@" BUILTIN_ACL_DOMAIN); + Server.AnnoyingClient->ModifyACL("/Root/PQ", FullTopicName, acl.SerializeAsString()); + + TReadRequest req; + req.MutableRead()->SetMaxCount(1); + if (!stream->Write(req)) { + ythrow yexception() << "write fail"; + } + + UNIT_ASSERT(stream->Read(&resp)); + UNIT_ASSERT(resp.HasError() && resp.GetError().GetCode() == NPersQueue::NErrorCode::EErrorCode::ACCESS_DENIED); + return; + } + Server.AnnoyingClient->GetClientInfo({FullTopicName}, clientId, true); + for (ui32 i = 0; i < 11; ++i) { + TReadResponse resp; + + UNIT_ASSERT(stream->Read(&resp)); + + if (useBatching) { + UNIT_ASSERT(resp.HasBatchedData()); + UNIT_ASSERT_VALUES_EQUAL(resp.GetBatchedData().PartitionDataSize(), 1); + UNIT_ASSERT_VALUES_EQUAL(resp.GetBatchedData().GetPartitionData(0).BatchSize(), 1); + UNIT_ASSERT_VALUES_EQUAL(resp.GetBatchedData().GetPartitionData(0).GetBatch(0).MessageDataSize(), 1); + UNIT_ASSERT_VALUES_EQUAL(resp.GetBatchedData().GetPartitionData(0).GetBatch(0).GetMessageData(0).GetOffset(), i); + } else { + UNIT_ASSERT(resp.HasData()); + UNIT_ASSERT_VALUES_EQUAL(resp.GetData().MessageBatchSize(), 1); + UNIT_ASSERT_VALUES_EQUAL(resp.GetData().GetMessageBatch(0).MessageSize(), 1); + UNIT_ASSERT_VALUES_EQUAL(resp.GetData().GetMessageBatch(0).GetMessage(0).GetOffset(), i); + } + } + //TODO: check here that read will never done UNIT_ASSERT(!stream->Read(&resp)); + { + for (ui32 i = 1; i < 11; ++i) { + TReadRequest req; + + req.MutableCommit()->AddCookie(i); + + if (!stream->Write(req)) { + ythrow yexception() << "write fail"; + } + } + ui32 i = 1; + while (i <= 10) { + TReadResponse resp; + + UNIT_ASSERT(stream->Read(&resp)); + Cerr << resp << "\n"; + UNIT_ASSERT(resp.HasCommit()); + UNIT_ASSERT(resp.GetCommit().CookieSize() > 0); + for (ui32 j = 0; j < resp.GetCommit().CookieSize(); ++j) { + UNIT_ASSERT( resp.GetCommit().GetCookie(j) == i); + ++i; + UNIT_ASSERT(i <= 11); + } + } + Server.AnnoyingClient->GetClientInfo({FullTopicName}, clientId, true); + } +} + +void TPQDataWriter::WaitWritePQServiceInitialization() { + TWriteRequest req; + TWriteResponse resp; + while (true) { + grpc::ClientContext context; + + auto stream = StubP_->WriteSession(&context); + UNIT_ASSERT(stream); + req.MutableInit()->SetTopic(ShortTopicName); + req.MutableInit()->SetSourceId("12345678"); + req.MutableInit()->SetProxyCookie(ProxyCookie); + + if (!stream->Write(req)) { + continue; + } + UNIT_ASSERT(stream->Read(&resp)); + if (resp.GetError().GetCode() == NPersQueue::NErrorCode::INITIALIZING) { + Sleep(TDuration::MilliSeconds(50)); + } else { + break; + } + } +} + +ui32 TPQDataWriter::InitSession(const TString& sourceId, ui32 pg, bool success, ui32 step) { + TWriteRequest req; + TWriteResponse resp; + + grpc::ClientContext context; + + auto stream = StubP_->WriteSession(&context); + + UNIT_ASSERT(stream); + req.MutableInit()->SetTopic(ShortTopicName); + req.MutableInit()->SetSourceId(sourceId); + req.MutableInit()->SetPartitionGroup(pg); + req.MutableInit()->SetProxyCookie(ProxyCookie); + + UNIT_ASSERT(stream->Write(req)); + UNIT_ASSERT(stream->Read(&resp)); + Cerr << "Init result: " << resp << "\n"; + //TODO: ensure topic creation - proxy already knows about new partitions, but tablet - no! + if (!success) { + UNIT_ASSERT(resp.HasError()); + return 0; + } else { + if (!resp.HasInit() && step < 5) { + Sleep(TDuration::MilliSeconds(100)); + return InitSession(sourceId, pg, success, step + 1); + } + UNIT_ASSERT(resp.HasInit()); + return resp.GetInit().GetPartition(); + } + return 0; +} + +ui32 TPQDataWriter::WriteImpl( + const TString& topic, const TVector<TString>& data, bool error, const TString& ticket, bool batch +) { + grpc::ClientContext context; + + auto stream = StubP_->WriteSession(&context); + UNIT_ASSERT(stream); + + // Send initial request. + TWriteRequest req; + TWriteResponse resp; + + req.MutableInit()->SetTopic(topic); + req.MutableInit()->SetSourceId(SourceId_); + req.MutableInit()->SetProxyCookie(ProxyCookie); + if (!ticket.empty()) + req.MutableCredentials()->SetTvmServiceTicket(ticket); + auto item = req.MutableInit()->MutableExtraFields()->AddItems(); + item->SetKey("key"); + item->SetValue("value"); + + if (!stream->Write(req)) { + ythrow yexception() << "write fail"; + } + + ui32 part = 0; + + UNIT_ASSERT(stream->Read(&resp)); + + if (!error) { + UNIT_ASSERT_C(resp.HasInit(), resp); + UNIT_ASSERT_C(!resp.GetInit().GetSessionId().empty(), resp); + part = resp.GetInit().GetPartition(); + } else { + Cerr << resp << "\n"; + UNIT_ASSERT(resp.HasError()); + return 0; + } + + // Send data requests. + Flush(data, stream, ticket, batch); + + Flush(data, stream, ticket, batch); + + Flush(data, stream, ticket, batch); + + Flush(data, stream, ticket, batch); + + //will cause only 4 answers in stream->Read - third call will fail, not blocks + stream->WritesDone(); + + UNIT_ASSERT(!stream->Read(&resp)); + + auto status = stream->Finish(); + UNIT_ASSERT(status.ok()); + return part; +} + +ui64 TPQDataWriter::ReadCookieFromMetadata(const std::multimap<grpc::string_ref, grpc::string_ref>& meta) const { + auto ci = meta.find("cookie"); + if (ci == meta.end()) { + ythrow yexception() << "server didn't provide the cookie"; + } else { + return FromString<ui64>(TStringBuf(ci->second.data(), ci->second.size())); + } +} + +void TPQDataWriter::InitializeChannel() { + Channel_ = grpc::CreateChannel("[::1]:" + ToString(Server.GrpcPort), grpc::InsecureChannelCredentials()); + Stub_ = NKikimrClient::TGRpcServer::NewStub(Channel_); + + grpc::ClientContext context; + NKikimrClient::TChooseProxyRequest request; + NKikimrClient::TResponse response; + auto status = Stub_->ChooseProxy(&context, request, &response); + UNIT_ASSERT(status.ok()); + Cerr << response << "\n"; + UNIT_ASSERT(response.GetStatus() == NMsgBusProxy::MSTATUS_OK); + ProxyCookie = response.GetProxyCookie(); + Channel_ = grpc::CreateChannel( + "[" + response.GetProxyName() + "]:" + ToString(Server.GrpcPort), + grpc::InsecureChannelCredentials() + ); + StubP_ = NPersQueue::PersQueueService::NewStub(Channel_); +} +} // namespace NKikimr::NPersQueueTests diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/data_writer.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/data_writer.h new file mode 100644 index 0000000000..c4e46ab08e --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/data_writer.h @@ -0,0 +1,117 @@ +#pragma once + +#include "test_server.h" + +#include <kikimr/yndx/api/grpc/persqueue.grpc.pb.h> + +#include <library/cpp/testing/unittest/registar.h> + +#include <grpc++/client_context.h> +#include <grpc++/create_channel.h> + +namespace NPersQueue::NTests { + +class TPQDataWriter { +public: + TPQDataWriter( + const TString& defaultTopicFullName, const TString& defaultShortTopicName, + const TString& sourceId, TTestServer& server + ) + : FullTopicName(defaultTopicFullName) + , ShortTopicName(defaultShortTopicName) + , SourceId_(sourceId) + , Server(server) + { + InitializeChannel(); + WaitWritePQServiceInitialization(); + } + + void Read(const TString& topic, const TString& clientId, const TString& ticket = "", bool error = false, + bool checkACL = false, bool useBatching = false, bool onlyCreate = false); + + void WaitWritePQServiceInitialization(); + + ui32 InitSession(const TString& sourceId, ui32 pg, bool success, ui32 step = 0); + + ui32 Write(const TString& topic, const TString& data, bool error = false, const TString& ticket = "") { + return WriteImpl(topic, {data}, error, ticket, false); + } + + ui32 WriteBatch(const TString& topic, const TVector<TString>& data, bool error = false, + const TString& ticket = "") { + return WriteImpl(topic, data, error, ticket, true); + } + +private: + ui32 WriteImpl(const TString& topic, const TVector<TString>& data, bool error, const TString& ticket, + bool batch); + + template<typename S> + void Flush(const TVector<TString>& data, S& stream, const TString& ticket, bool batch) { + TWriteRequest request; + TWriteResponse response; + + TVector<ui64> allSeqNo; + if (batch) { + for (const TString& d: data) { + ui64 seqNo = AtomicIncrement(SeqNo_); + allSeqNo.push_back(seqNo); + auto *mutableData = request.MutableDataBatch()->AddData(); + mutableData->SetSeqNo(seqNo); + mutableData->SetData(d); + } + } else { + ui64 seqNo = AtomicIncrement(SeqNo_); + allSeqNo.push_back(seqNo); + request.MutableData()->SetSeqNo(seqNo); + request.MutableData()->SetData(JoinSeq("\n", data)); + } + if (!ticket.empty()) { + request.MutableCredentials()->SetTvmServiceTicket(ticket); + } + + Cerr << "request: " << request << Endl; + if (!stream->Write(request)) { + ythrow yexception() << "write fail"; + } + + UNIT_ASSERT(stream->Read(&response)); + if (batch) { + UNIT_ASSERT_C(response.HasAckBatch(), response); + UNIT_ASSERT_VALUES_EQUAL(data.size(), response.GetAckBatch().AckSize()); + for (size_t i = 0; i < data.size(); ++i) { + const auto& ack = response.GetAckBatch().GetAck(i); + UNIT_ASSERT(!ack.GetAlreadyWritten()); + UNIT_ASSERT(!ack.HasStat()); + UNIT_ASSERT_VALUES_EQUAL(ack.GetSeqNo(), allSeqNo[i]); + } + UNIT_ASSERT(response.GetAckBatch().HasStat()); + } else { + const auto& ack = response.GetAck(); + UNIT_ASSERT(!ack.GetAlreadyWritten()); + UNIT_ASSERT_VALUES_EQUAL(ack.GetSeqNo(), allSeqNo[0]); + UNIT_ASSERT(ack.HasStat()); + } + } + + ui64 ReadCookieFromMetadata(const std::multimap<grpc::string_ref, grpc::string_ref>& meta) const; + + void InitializeChannel(); + +private: + TString FullTopicName; + TString ShortTopicName; + const TString SourceId_; + + TTestServer& Server; + + TAtomic SeqNo_ = 1; + + //! Сетевой канал взаимодействия с proxy-сервером. + std::shared_ptr<grpc::Channel> Channel_; + std::unique_ptr<NKikimrClient::TGRpcServer::Stub> Stub_; + std::unique_ptr<PersQueueService::Stub> StubP_; + + ui64 ProxyCookie = 0; +}; +} // NKikimr::NPersQueueTests diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/sdk_test_setup.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/sdk_test_setup.h new file mode 100644 index 0000000000..9f0712933a --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/sdk_test_setup.h @@ -0,0 +1,323 @@ +#pragma once +#include "test_server.h" +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> +#include <ydb/library/persqueue/topic_parser_public/topic_parser.h> +#include <library/cpp/logger/log.h> +#include <util/system/tempfile.h> + +#define TEST_CASE_NAME (this->Name_) + +namespace NPersQueue { +class SDKTestSetup { +protected: + TString TestCaseName; + + THolder<TTempFileHandle> NetDataFile; + THashMap<TString, NKikimr::NPersQueueTests::TPQTestClusterInfo> DataCenters; + TString LocalDC = "dc1"; + TTestServer Server = TTestServer(false /* don't start */); + + TLog Log = TLog("cerr"); + + TPQLibSettings PQLibSettings; + THolder<TPQLib> PQLib; + +public: + SDKTestSetup(const TString& testCaseName, bool start = true) + : TestCaseName(testCaseName) + { + InitOptions(); + if (start) { + Start(); + } + } + + void InitOptions() { + Log.SetFormatter([testCaseName = TestCaseName](ELogPriority priority, TStringBuf message) { + return TStringBuilder() << TInstant::Now() << " :" << testCaseName << " " << priority << ": " << message << Endl; + }); + Server.GrpcServerOptions.SetGRpcShutdownDeadline(TDuration::Max()); + // Default TTestServer value for 'MaxReadCookies' is 10. With this value the tests are flapping with two errors: + // 1. 'got more than 10 unordered cookies to commit 12' + // 2. 'got more than 10 uncommitted reads' + Server.ServerSettings.PQConfig.Clear(); + Server.ServerSettings.PQConfig.SetEnabled(true); + Server.ServerSettings.PQConfig.SetRemoteClusterEnabledDelaySec(1); + Server.ServerSettings.PQConfig.SetCloseClientSessionWithEnabledRemotePreferredClusterDelaySec(1); + Server.ServerSettings.PQClusterDiscoveryConfig.SetEnabled(true); + SetNetDataViaFile("::1/128\t" + GetLocalCluster()); + + auto seed = TInstant::Now().MicroSeconds(); + // This makes failing randomized tests (for example with NUnitTest::RandomString(size, std::rand()) calls) reproducable + Log << TLOG_INFO << "Random seed for debugging is " << seed; + std::srand(seed); + } + + void Start(bool waitInit = true, bool addBrokenDatacenter = !std::getenv("PERSQUEUE_GRPC_API_V1_ENABLED")) { + Server.StartServer(false); + //Server.EnableLogs({NKikimrServices::PQ_WRITE_PROXY, NKikimrServices::PQ_READ_PROXY}); + Server.AnnoyingClient->InitRoot(); + if (DataCenters.empty()) { + THashMap<TString, NKikimr::NPersQueueTests::TPQTestClusterInfo> dataCenters; + dataCenters.emplace("dc1", NKikimr::NPersQueueTests::TPQTestClusterInfo{TStringBuilder() << "localhost:" << Server.GrpcPort, true}); + if (addBrokenDatacenter) { + dataCenters.emplace("dc2", NKikimr::NPersQueueTests::TPQTestClusterInfo{"dc2.logbroker.yandex.net", false}); + } + Server.AnnoyingClient->InitDCs(dataCenters); + } else { + Server.AnnoyingClient->InitDCs(DataCenters, LocalDC); + } + Server.AnnoyingClient->InitSourceIds(); + CreateTopic(GetTestTopic(), GetLocalCluster()); + if (waitInit) { + Server.WaitInit(GetTestTopic()); + } + PQLibSettings.DefaultLogger = MakeIntrusive<TCerrLogger>(TLOG_DEBUG); + PQLib = MakeHolder<TPQLib>(PQLibSettings); + } + + THolder<TPQLib>& GetPQLib() { + return PQLib; + } + + const TPQLibSettings& GetPQLibSettings() const { + return PQLibSettings; + } + + TString GetTestTopic() const { + return "topic1"; + } + + TString GetTestClient() const { + return "test-reader"; + } + + TString GetTestMessageGroupId() const { + return "test-message-group-id"; + } + + TString GetLocalCluster() const { + return LocalDC; + } + + + NGrpc::TServerOptions& GetGrpcServerOptions() { + return Server.GrpcServerOptions; + } + + void SetNetDataViaFile(const TString& netDataTsv) { + NetDataFile = MakeHolder<TTempFileHandle>("netData.tsv"); + NetDataFile->Write(netDataTsv.Data(), netDataTsv.Size()); + NetDataFile->FlushData(); + Server.ServerSettings.NetClassifierConfig.SetNetDataFilePath(NetDataFile->Name()); + } + + TProducerSettings GetProducerSettings() const { + TProducerSettings producerSettings; + producerSettings.Topic = GetTestTopic(); + producerSettings.SourceId = GetTestMessageGroupId(); + producerSettings.Server = TServerSetting{"localhost", Server.GrpcPort}; + producerSettings.Server.Database = "/Root"; + return producerSettings; + } + + TConsumerSettings GetConsumerSettings() const { + TConsumerSettings consumerSettings; + consumerSettings.Topics = {GetTestTopic()}; + consumerSettings.ClientId = GetTestClient(); + consumerSettings.Server = TServerSetting{"localhost", Server.GrpcPort}; + consumerSettings.UseV2RetryPolicyInCompatMode = true; + return consumerSettings; + } + + TLog& GetLog() { + return Log; + } + + + template <class TConsumerOrProducer> + void Start(const THolder<TConsumerOrProducer>& obj) { + auto startFuture = obj->Start(); + const auto& initResponse = startFuture.GetValueSync(); + UNIT_ASSERT_C(!initResponse.Response.HasError(), "Failed to start: " << initResponse.Response); + } + + THolder<IConsumer> StartConsumer(const TConsumerSettings& settings) { + THolder<IConsumer> consumer = GetPQLib()->CreateConsumer(settings); + Start(consumer); + return consumer; + } + + THolder<IConsumer> StartConsumer() { + return StartConsumer(GetConsumerSettings()); + } + + THolder<IProducer> StartProducer(const TProducerSettings& settings) { + THolder<IProducer> producer = GetPQLib()->CreateProducer(settings); + Start(producer); + return producer; + } + + THolder<IProducer> StartProducer() { + return StartProducer(GetProducerSettings()); + } + + void WriteToTopic(const TVector<TString>& data, IProducer* producer = nullptr) { + THolder<IProducer> localProducer; + if (!producer) { + localProducer = StartProducer(); + producer = localProducer.Get(); + } + TVector<NThreading::TFuture<TProducerCommitResponse>> resps; + for (const TString& d : data) { + Log << TLOG_INFO << "WriteToTopic: " << d; + resps.push_back(producer->Write(d)); + } + for (NThreading::TFuture<TProducerCommitResponse>& r : resps) { + UNIT_ASSERT_C(!r.GetValueSync().Response.HasError(), r.GetValueSync().Response.GetError()); + } + } + + // Read set of sequences from topic + void ReadFromTopic(const TVector<TVector<TString>>& data, bool commit = true, IConsumer* consumer = nullptr) { + THolder<IConsumer> localConsumer; + if (!consumer) { + localConsumer = StartConsumer(); + consumer = localConsumer.Get(); + } + TVector<size_t> positions(data.size()); // Initialy zeroes. + + int wholeCount = 0; + for (const TVector<TString>& seq : data) { + wholeCount += seq.size(); + } + + TSet<ui64> cookies; + + auto processCommit = [&](const TConsumerMessage& resp) { + Log << TLOG_INFO << "ReadFromTopic. Committed: " << resp.Response.GetCommit(); + for (ui64 cookie : resp.Response.GetCommit().GetCookie()) { + UNIT_ASSERT_VALUES_EQUAL(cookies.erase(cookie), 1); + } + }; + + while (wholeCount > 0) { + auto event = consumer->GetNextMessage(); + const auto& resp = event.GetValueSync(); + UNIT_ASSERT_C(!resp.Response.HasError(), resp.Response); + if (!resp.Response.HasData()) { + if (resp.Response.HasCommit()) { + processCommit(resp); + } + continue; + } + Log << TLOG_INFO << "ReadFromTopic. Data: " << resp.Response.GetData(); + UNIT_ASSERT(cookies.insert(resp.Response.GetData().GetCookie()).second); + for (const auto& batch : resp.Response.GetData().GetMessageBatch()) { + // find proper sequence + const TString& firstData = batch.GetMessage(0).GetData(); + size_t seqIndex = 0; + for (; seqIndex < positions.size(); ++seqIndex) { + if (positions[seqIndex] >= data[seqIndex].size()) { // Already seen. + continue; + } + size_t& seqPos = positions[seqIndex]; + const TString& expectedData = data[seqIndex][seqPos]; + if (expectedData == firstData) { + UNIT_ASSERT(batch.MessageSize() <= data[seqIndex].size() - positions[seqIndex]); + ++seqPos; + --wholeCount; + // Match. + for (size_t msgIndex = 1; msgIndex < batch.MessageSize(); ++msgIndex, ++seqPos, --wholeCount) { + UNIT_ASSERT_STRINGS_EQUAL(batch.GetMessage(msgIndex).GetData(), data[seqIndex][seqPos]); + } + break; + } + } + UNIT_ASSERT_LT_C(seqIndex, positions.size(), resp.Response); + } + + if (commit) { + consumer->Commit({resp.Response.GetData().GetCookie()}); + } + } + while (commit && !cookies.empty()) { + auto event = consumer->GetNextMessage(); + const auto& resp = event.GetValueSync(); + UNIT_ASSERT_C(!resp.Response.HasError(), resp.Response); + if (!resp.Response.HasCommit()) { + continue; + } + processCommit(resp); + } + UNIT_ASSERT_VALUES_EQUAL(wholeCount, 0); + } + + void SetSingleDataCenter(const TString& name = "dc1") { + UNIT_ASSERT( + DataCenters.insert(std::make_pair( + name, + NKikimr::NPersQueueTests::TPQTestClusterInfo{TStringBuilder() << "localhost:" << Server.GrpcPort, true} + )).second + ); + LocalDC = name; + } + + void AddDataCenter(const TString& name, const TString& address, bool enabled = true, bool setSelfAsDc = true) { + if (DataCenters.empty() && setSelfAsDc) { + SetSingleDataCenter(); + } + NKikimr::NPersQueueTests::TPQTestClusterInfo info{ + address, + enabled + }; + UNIT_ASSERT(DataCenters.insert(std::make_pair(name, info)).second); + } + + void AddDataCenter(const TString& name, const SDKTestSetup& cluster, bool enabled = true, bool setSelfAsDc = true) { + AddDataCenter(name, TStringBuilder() << "localhost:" << cluster.Server.GrpcPort, enabled, setSelfAsDc); + } + + void EnableDataCenter(const TString& name) { + auto iter = DataCenters.find(name); + UNIT_ASSERT(iter != DataCenters.end()); + Server.AnnoyingClient->UpdateDcEnabled(name, true); + } + void DisableDataCenter(const TString& name) { + auto iter = DataCenters.find(name); + UNIT_ASSERT(iter != DataCenters.end()); + Server.AnnoyingClient->UpdateDcEnabled(name, false); + } + + void ShutdownGRpc() { + Server.ShutdownGRpc(); + } + + void EnableGRpc() { + Server.EnableGRpc(); + Server.WaitInit(GetTestTopic()); + } + + void KickTablets() { + for (ui32 i = 0; i < Server.CleverServer->StaticNodes(); i++) { + Server.AnnoyingClient->MarkNodeInHive(Server.CleverServer->GetRuntime(), i, false); + Server.AnnoyingClient->KickNodeInHive(Server.CleverServer->GetRuntime(), i); + } + } + void AllowTablets() { + for (ui32 i = 0; i < Server.CleverServer->StaticNodes(); i++) { + Server.AnnoyingClient->MarkNodeInHive(Server.CleverServer->GetRuntime(), i, true); + } + } + + void CreateTopic(const TString& topic, const TString& cluster, size_t partitionsCount = 1) { + Server.AnnoyingClient->CreateTopicNoLegacy(BuildFullTopicName(topic, cluster), partitionsCount); + } + + void KillPqrb(const TString& topic, const TString& cluster) { + auto describeResult = Server.AnnoyingClient->Ls(TStringBuilder() << "/Root/PQ/" << BuildFullTopicName(topic, cluster)); + UNIT_ASSERT_C(describeResult->Record.GetPathDescription().HasPersQueueGroup(), describeResult->Record); + Server.AnnoyingClient->KillTablet(*Server.CleverServer, describeResult->Record.GetPathDescription().GetPersQueueGroup().GetBalancerTabletID()); + } +}; +} diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_pqlib.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_pqlib.cpp new file mode 100644 index 0000000000..f5d0b63cb3 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_pqlib.cpp @@ -0,0 +1,95 @@ +#include "test_pqlib.h" + +namespace NPersQueue::NTests { + + std::tuple<THolder<IProducer>, TProducerCreateResponse> TTestPQLib::CreateProducer( + const TProducerSettings& settings, bool deprecated + ) { + auto producer = PQLib->CreateProducer(settings, Logger, deprecated); + TProducerCreateResponse response = producer->Start().GetValueSync(); + return std::tuple<THolder<IProducer>, TProducerCreateResponse>(std::move(producer), response); + } + + std::tuple<THolder<IProducer>, TProducerCreateResponse> TTestPQLib::CreateProducer( + const TString& topic, + const TString& sourceId, + std::optional<ui32> partitionGroup, + std::optional<NPersQueueCommon::ECodec> codec, + std::optional<bool> reconnectOnFailure, + bool deprecated + ) { + TProducerSettings s; + s.Server = TServerSetting{"localhost", Server.GrpcPort}; + s.Topic = topic; + s.SourceId = sourceId; + if (partitionGroup) { + s.PartitionGroup = partitionGroup.value(); + } + if (codec) { + s.Codec = codec.value(); + } + if (reconnectOnFailure) { + s.ReconnectOnFailure = reconnectOnFailure.value(); + } + return CreateProducer(s, deprecated); + } + + std::tuple<THolder<IConsumer>, TConsumerCreateResponse> TTestPQLib::CreateConsumer( + const TConsumerSettings& settings + ) { + auto consumer = PQLib->CreateConsumer(settings, Logger); + TConsumerCreateResponse response = consumer->Start().GetValueSync(); + return std::tuple<THolder<IConsumer>, TConsumerCreateResponse>(std::move(consumer), response); + } + + std::tuple<THolder<IConsumer>, TConsumerCreateResponse> TTestPQLib::CreateDeprecatedConsumer( + const TConsumerSettings& settings + ) { + auto consumer = PQLib->CreateConsumer(settings, Logger, true); + TConsumerCreateResponse response = consumer->Start().GetValueSync(); + return std::tuple<THolder<IConsumer>, TConsumerCreateResponse>(std::move(consumer), response); + } + + TConsumerSettings TTestPQLib::MakeSettings( + const TVector<TString>& topics, + const TString& consumerName + ) { + TConsumerSettings s; + s.Server = TServerSetting{"localhost", Server.GrpcPort}; + s.ClientId = consumerName; + s.Topics = topics; + return s; + } + + std::tuple<THolder<IConsumer>, TConsumerCreateResponse> TTestPQLib::CreateConsumer( + const TVector<TString>& topics, + const TString& consumerName, + std::optional<ui32> maxCount, + std::optional<bool> useLockSession, + std::optional<bool> readMirroredPartitions, + std::optional<bool> unpack, + std::optional<ui32> maxInflyRequests, + std::optional<ui32> maxMemoryUsage + ) { + auto s = MakeSettings(topics, consumerName); + if (maxCount) { + s.MaxCount = maxCount.value(); + } + if (unpack) { + s.Unpack = unpack.value(); + } + if (readMirroredPartitions) { + s.ReadMirroredPartitions = readMirroredPartitions.value(); + } + if (maxInflyRequests) { + s.MaxInflyRequests = maxInflyRequests.value(); + } + if (maxMemoryUsage) { + s.MaxMemoryUsage = maxMemoryUsage.value(); + } + if (useLockSession) { + s.UseLockSession = useLockSession.value(); + } + return CreateConsumer(s); + } +} // namespace NPersQueue::NTests
\ No newline at end of file diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_pqlib.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_pqlib.h new file mode 100644 index 0000000000..8e6fb6832d --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_pqlib.h @@ -0,0 +1,43 @@ +#pragma once +#include "test_server.h" + +namespace NPersQueue::NTests { + +class TTestPQLib { +public: + TTestPQLib(TTestServer& server) + : Logger(new TCerrLogger(DEBUG_LOG_LEVEL)) + , Server(server) + { + TPQLibSettings settings; + settings.DefaultLogger = Logger; + PQLib = MakeHolder<TPQLib>(settings); + } + + std::tuple<THolder<IProducer>, TProducerCreateResponse> CreateProducer( + const TProducerSettings& settings, bool deprecated + ); + + std::tuple<THolder<IProducer>, TProducerCreateResponse> CreateProducer( + const TString& topic, const TString& sourceId, std::optional<ui32> partitionGroup = {}, + std::optional<NPersQueueCommon::ECodec> codec = ECodec::RAW, std::optional<bool> reconnectOnFailure = {}, + bool deprecated = false); + + std::tuple<THolder<IConsumer>, TConsumerCreateResponse> CreateConsumer(const TConsumerSettings& settings); + + std::tuple<THolder<IConsumer>, TConsumerCreateResponse> CreateDeprecatedConsumer(const TConsumerSettings& settings); + + TConsumerSettings MakeSettings(const TVector<TString>& topics, const TString& consumerName = "user"); + + std::tuple<THolder<IConsumer>, TConsumerCreateResponse> CreateConsumer( + const TVector<TString>& topics, const TString& consumerName = "user", std::optional<ui32> maxCount = {}, + std::optional<bool> useLockSession = {}, std::optional<bool> readMirroredPartitions = {}, + std::optional<bool> unpack = {}, std::optional<ui32> maxInflyRequests = {}, + std::optional<ui32> maxMemoryUsage = {}); + +private: + TIntrusivePtr<TCerrLogger> Logger; + TTestServer& Server; + THolder<TPQLib> PQLib; +}; +} // namespace NPersQueue::NTests
\ No newline at end of file diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.cpp b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.cpp new file mode 100644 index 0000000000..6cd1e8094a --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.cpp @@ -0,0 +1,35 @@ +#include "test_server.h" + +#include <ydb/core/client/server/msgbus_server_pq_metacache.h> +#include <kikimr/yndx/grpc_services/persqueue/persqueue.h> +#include <kikimr/yndx/persqueue/msgbus_server/read_session_info.h> + +namespace NPersQueue { + +const TVector<NKikimrServices::EServiceKikimr> TTestServer::LOGGED_SERVICES = { + NKikimrServices::PERSQUEUE, + NKikimrServices::PQ_METACACHE, + NKikimrServices::PQ_READ_PROXY, + NKikimrServices::PQ_WRITE_PROXY, + NKikimrServices::PQ_MIRRORER, +}; + +void TTestServer::PatchServerSettings() { + ServerSettings.RegisterGrpcService<NKikimr::NGRpcService::TGRpcPersQueueService>( + "pq", + NKikimr::NMsgBusProxy::CreatePersQueueMetaCacheV2Id() + ); + ServerSettings.SetPersQueueGetReadSessionsInfoWorkerFactory( + std::make_shared<NKikimr::NMsgBusProxy::TPersQueueGetReadSessionsInfoWorkerWithPQv0Factory>() + ); +} + +void TTestServer::StartIfNeeded(bool start) { + if (start) { + StartServer(); + } +} + + + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h new file mode 100644 index 0000000000..090d9aa47f --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h @@ -0,0 +1,129 @@ +#pragma once +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> + +#include <ydb/core/testlib/test_pq_client.h> +#include <library/cpp/grpc/server/grpc_server.h> + +#include <library/cpp/testing/unittest/registar.h> +#include <library/cpp/testing/unittest/tests_data.h> + +#include <util/system/tempfile.h> + +namespace NPersQueue { + +static constexpr int DEBUG_LOG_LEVEL = 7; + +class TTestServer { +public: + TTestServer(bool start = true, TMaybe<TSimpleSharedPtr<TPortManager>> portManager = Nothing()) + : PortManager(portManager.GetOrElse(MakeSimpleShared<TPortManager>())) + , Port(PortManager->GetPort(2134)) + , GrpcPort(PortManager->GetPort(2135)) + , ServerSettings(NKikimr::NPersQueueTests::PQSettings(Port).SetGrpcPort(GrpcPort)) + , GrpcServerOptions(NGrpc::TServerOptions().SetHost("[::1]").SetPort(GrpcPort)) + , PQLibSettings(TPQLibSettings{ .DefaultLogger = new TCerrLogger(DEBUG_LOG_LEVEL) }) + , PQLib(new TPQLib(PQLibSettings)) + { + PatchServerSettings(); + StartIfNeeded(start); + } + + TTestServer(const NKikimr::Tests::TServerSettings& settings, bool start = true) + : PortManager(MakeSimpleShared<TPortManager>()) + , Port(PortManager->GetPort(2134)) + , GrpcPort(PortManager->GetPort(2135)) + , ServerSettings(settings) + , GrpcServerOptions(NGrpc::TServerOptions().SetHost("[::1]").SetPort(GrpcPort)) + , PQLibSettings(TPQLibSettings{ .DefaultLogger = new TCerrLogger(DEBUG_LOG_LEVEL) }) + , PQLib(new TPQLib(PQLibSettings)) + { + ServerSettings.Port = Port; + ServerSettings.SetGrpcPort(GrpcPort); + PatchServerSettings(); + StartIfNeeded(start); + } + + void StartServer(bool doClientInit = true) { + PrepareNetDataFile(); + CleverServer = MakeHolder<NKikimr::Tests::TServer>(ServerSettings); + CleverServer->EnableGRpc(GrpcServerOptions); + AnnoyingClient = MakeHolder<NKikimr::NPersQueueTests::TFlatMsgBusPQClient>(ServerSettings, GrpcPort); + EnableLogs(LOGGED_SERVICES); + if (doClientInit) { + AnnoyingClient->FullInit(); + } + } + + void ShutdownGRpc() { + CleverServer->ShutdownGRpc(); + } + + void EnableGRpc() { + CleverServer->EnableGRpc(GrpcServerOptions); + } + + void ShutdownServer() { + CleverServer = nullptr; + } + + void RestartServer() { + ShutdownServer(); + StartServer(); + } + + void EnableLogs(const TVector<NKikimrServices::EServiceKikimr> services, + NActors::NLog::EPriority prio = NActors::NLog::PRI_DEBUG) { + Y_VERIFY(CleverServer != nullptr, "Start server before enabling logs"); + for (auto s : services) { + CleverServer->GetRuntime()->SetLogPriority(s, prio); + } + } + + void WaitInit(const TString& topic) { + TProducerSettings s; + s.Topic = topic; + s.Server = TServerSetting{"localhost", GrpcPort}; + s.SourceId = "src"; + + while (PQLib->CreateProducer(s, {}, false)->Start().GetValueSync().Response.HasError()) { + Sleep(TDuration::MilliSeconds(200)); + } + } + + bool PrepareNetDataFile(const TString& content = "::1/128\tdc1") { + if (NetDataFile) + return false; + NetDataFile = MakeHolder<TTempFileHandle>("netData.tsv"); + NetDataFile->Write(content.Data(), content.Size()); + NetDataFile->FlushData(); + ServerSettings.NetClassifierConfig.SetNetDataFilePath(NetDataFile->Name()); + return true; + } + + void UpdateDC(const TString& name, bool local, bool enabled) { + AnnoyingClient->UpdateDC(name, local, enabled); + } + +private: + void PatchServerSettings(); + void StartIfNeeded(bool start); + +public: + TSimpleSharedPtr<TPortManager> PortManager; + ui16 Port; + ui16 GrpcPort; + + THolder<NKikimr::Tests::TServer> CleverServer; + NKikimr::Tests::TServerSettings ServerSettings; + NGrpc::TServerOptions GrpcServerOptions; + THolder<TTempFileHandle> NetDataFile; + + THolder<NKikimr::NPersQueueTests::TFlatMsgBusPQClient> AnnoyingClient; + + TPQLibSettings PQLibSettings; + THolder<TPQLib> PQLib; + + static const TVector<NKikimrServices::EServiceKikimr> LOGGED_SERVICES; +}; + +} // namespace NPersQueue diff --git a/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h new file mode 100644 index 0000000000..13b4da2772 --- /dev/null +++ b/kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_utils.h @@ -0,0 +1,101 @@ +#pragma once +#include <util/generic/ptr.h> +#include <util/generic/size_literals.h> +#include <library/cpp/threading/chunk_queue/queue.h> +#include <util/generic/overloaded.h> +#include <library/cpp/testing/unittest/registar.h> + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/sdk_test_setup.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> + +namespace NPersQueue { + +using namespace NThreading; +using namespace NYdb::NPersQueue; +using namespace NKikimr; +using namespace NKikimr::NPersQueueTests; +//using namespace NPersQueue::V1; + +template <class TPQLibObject> +void DestroyAndWait(THolder<TPQLibObject>& object) { + if (object) { + auto isDead = object->IsDead(); + object = nullptr; + isDead.GetValueSync(); + } +} + +inline bool GrpcV1EnabledByDefault() { + static const bool enabled = std::getenv("PERSQUEUE_GRPC_API_V1_ENABLED"); + return enabled; +} + +class TCallbackCredentialsProvider : public ICredentialsProvider { + std::function<void(NPersQueue::TCredentials*)> Callback; +public: + TCallbackCredentialsProvider(std::function<void(NPersQueue::TCredentials*)> callback) + : Callback(std::move(callback)) + {} + + void FillAuthInfo(NPersQueue::TCredentials* authInfo) const { + Callback(authInfo); + } +}; + +struct TWriteResult { + bool Ok = false; + // No acknowledgement is expected from a writer under test + bool NoWait = false; + TString ResponseDebugString = TString(); +}; + +struct TAcknowledgableMessage { + TString Value; + ui64 SequenceNumber; + TInstant CreatedAt; + TPromise<TWriteResult> AckPromise; +}; + +class IClientEventLoop { +protected: + std::atomic_bool MayStop; + std::atomic_bool MustStop; + bool Stopped = false; + std::unique_ptr<TThread> Thread; + TLog Log; + +public: + IClientEventLoop() + : MayStop() + , MustStop() + , MessageBuffer() + {} + + void AllowStop() { + MayStop = true; + } + + void WaitForStop() { + if (!Stopped) { + Log << TLOG_INFO << "Wait for writer to die on itself"; + Thread->Join(); + Log << TLOG_INFO << "Client write event loop stopped"; + } + Stopped = true; + } + + virtual ~IClientEventLoop() { + MustStop = true; + if (!Stopped) { + Log << TLOG_INFO << "Wait for client write event loop to stop"; + Thread->Join(); + Log << TLOG_INFO << "Client write event loop stopped"; + } + Stopped = true; + } + + TManyOneQueue<TAcknowledgableMessage> MessageBuffer; + +}; + +} // namespace NPersQueue diff --git a/kikimr/public/README.md b/kikimr/public/README.md new file mode 100644 index 0000000000..b7422d385d --- /dev/null +++ b/kikimr/public/README.md @@ -0,0 +1,6 @@ +# Public APIs and Client SDKs for Yandex Database (YDB) + +`api` -- public API for YDB is implemented via gRPC framework; +`lib` -- some common libraries used by public SDKs and YDB core itself; +`sdk` -- programming language specific client SDK. If you are building an application on top of YDB these SDKs are for you; +`tools` -- developer tools: recipe to be used in tests, program to setup local YDB cluster. diff --git a/kikimr/public/sdk/cpp/README.md b/kikimr/public/sdk/cpp/README.md new file mode 100644 index 0000000000..adc7f4e4b5 --- /dev/null +++ b/kikimr/public/sdk/cpp/README.md @@ -0,0 +1,3 @@ +C++ SDK for Yandex Database (YDB) +--- +Officially supported C++ client for YDB. diff --git a/kikimr/public/sdk/cpp/client/CHANGELOG.md b/kikimr/public/sdk/cpp/client/CHANGELOG.md new file mode 100644 index 0000000000..bd8515352a --- /dev/null +++ b/kikimr/public/sdk/cpp/client/CHANGELOG.md @@ -0,0 +1,60 @@ +## 2.2.0 ## + +* Disable client query cache by default + +## 2.1.0 ## + +* Allow c++ sdk to create sessions in different ydb servers. + +## 2.0.0 ## + +* Remove request migrator feature from c++ sdk. + +## 1.6.2 ## + +* Fix SSL settings override for client + +## 1.6.1 ## + +* Do not wrap status in to CLIENT_DISCOVERY_FAILED in case of discovery error + +## 1.6.0 ## + +* Experimental interface to execute query from credential provider + +## 1.5.1 ## + +* Do not miss exceptions thrown by lambdas given to RetryOperation() method. + +## 1.5.0 ## + +* Added support of session graceful shutdown protocol + +## 1.4.0 ## + +* Support for socket timeout settings + +## 1.3.0 ## + +* Added user account key support for iam authentication + +## 1.2.0 ## + +* Added creation from environment support for DriverConfig + +## 1.1.0 ## + +* Support for grpc channel buffer limit settings + +## 1.0.2 ## + +* Added connection string support for CommonClientSettings, DriverConfig + +## 1.0.1 ## + +* Try to bind endpoint to session for + BeginTransaction, CommitTransaction, RollbackTransaction, KeepAlive, CloseSession, ExplainDataQuery requests + +## 1.0.0 ## + +* Start initial changelog. diff --git a/kikimr/public/sdk/cpp/client/iam/CMakeLists.txt b/kikimr/public/sdk/cpp/client/iam/CMakeLists.txt new file mode 100644 index 0000000000..27a57988b5 --- /dev/null +++ b/kikimr/public/sdk/cpp/client/iam/CMakeLists.txt @@ -0,0 +1,25 @@ + +# This file was gererated by the build system used internally in the Yandex monorepo. +# Only simple modifications are allowed (adding source-files to targets, adding simple properties +# like target_include_directories). These modifications will be ported to original +# ya.make files by maintainers. Any complex modifications which can't be ported back to the +# original buildsystem will not be accepted. + + + +add_library(cpp-client-iam) +target_link_libraries(cpp-client-iam PUBLIC + contrib-libs-cxxsupp + yutil + cloud-iam-v1 + priv-iam-v1 + cpp-grpc-client + cpp-http-simple + library-cpp-json + cpp-threading-atomic + public-lib-jwt + client-ydb_types-credentials +) +target_sources(cpp-client-iam PRIVATE + ${CMAKE_SOURCE_DIR}/kikimr/public/sdk/cpp/client/iam/iam.cpp +) diff --git a/kikimr/public/sdk/cpp/client/iam/iam.cpp b/kikimr/public/sdk/cpp/client/iam/iam.cpp new file mode 100644 index 0000000000..ff8f79d4fd --- /dev/null +++ b/kikimr/public/sdk/cpp/client/iam/iam.cpp @@ -0,0 +1,361 @@ +#include "iam.h" + +#include <cloud/bitbucket/public-api/yandex/cloud/iam/v1/iam_token_service.pb.h> +#include <cloud/bitbucket/public-api/yandex/cloud/iam/v1/iam_token_service.grpc.pb.h> + +#include <cloud/bitbucket/private-api/yandex/cloud/priv/iam/v1/iam_token_service.pb.h> +#include <cloud/bitbucket/private-api/yandex/cloud/priv/iam/v1/iam_token_service.grpc.pb.h> + + +#include <library/cpp/grpc/client/grpc_client_low.h> +#include <library/cpp/threading/atomic/bool.h> +#include <library/cpp/threading/future/core/future.h> +#include <library/cpp/json/json_reader.h> +#include <library/cpp/http/simple/http_client.h> + +#include <util/system/spinlock.h> +#include <util/stream/file.h> +#include <util/string/builder.h> + +#include <chrono> + +using namespace yandex::cloud::iam::v1; + + +using namespace NGrpc; + +namespace NYdb { + +constexpr TDuration BACKOFF_START = TDuration::MilliSeconds(50); +constexpr TDuration BACKOFF_MAX = TDuration::Seconds(10); + +class TIAMCredentialsProvider : public ICredentialsProvider { +public: + TIAMCredentialsProvider(const TIamHost& params) + : HttpClient_(TSimpleHttpClient(params.Host, params.Port)) + , Request_("/computeMetadata/v1/instance/service-accounts/default/token") + , NextTicketUpdate_(TInstant::Zero()) + , RefreshPeriod_(params.RefreshPeriod) + { + GetTicket(); + } + + TStringType GetAuthInfo() const override { + if (TInstant::Now() >= NextTicketUpdate_) { + GetTicket(); + } + return Ticket_; + } + + bool IsValid() const override { + return true; + } + +private: + TSimpleHttpClient HttpClient_; + TStringType Request_; + mutable TStringType Ticket_; + mutable TInstant NextTicketUpdate_; + TDuration RefreshPeriod_; + + void GetTicket() const { + try { + TStringStream out; + TSimpleHttpClient::THeaders headers; + headers["Metadata-Flavor"] = "Google"; + HttpClient_.DoGet(Request_, &out, headers); + NJson::TJsonValue resp; + NJson::ReadJsonTree(&out, &resp, true); + + auto respMap = resp.GetMap(); + + if (auto it = respMap.find("access_token"); it == respMap.end()) + ythrow yexception() << "Result doesn't contain access_token"; + else if (TString ticket = it->second.GetStringSafe(); ticket.empty()) + ythrow yexception() << "Got empty ticket"; + else + Ticket_ = std::move(ticket); + + if (auto it = respMap.find("expires_in"); it == respMap.end()) + ythrow yexception() << "Result doesn't contain expires_in"; + else { + const TDuration expiresIn = TDuration::Seconds(it->second.GetUInteger()); + + NextTicketUpdate_ = TInstant::Now() + std::max(expiresIn, RefreshPeriod_); + } + } catch (...) { + } + } +}; + + +template<typename TRequest, typename TResponse, typename TService> +class TGrpcIamCredentialsProvider : public ICredentialsProvider { +protected: + using TRequestFiller = std::function<void(TRequest&)>; + +private: + class TImpl : public std::enable_shared_from_this<TGrpcIamCredentialsProvider<TRequest, TResponse, TService>::TImpl> { + public: + TImpl(const TIamEndpoint& iamEndpoint, const TRequestFiller& requestFiller) + : Client(MakeHolder<NGrpc::TGRpcClientLow>()) + , Connection_(nullptr) + , Ticket_("") + , NextTicketUpdate_(TInstant::Zero()) + , IamEndpoint_(iamEndpoint) + , RequestFiller_(requestFiller) + , RequestInflight_(false) + , LastRequestError_("") + , NeedStop_(false) + , BackoffTimeout_(BACKOFF_START) + , Lock_() + { + TGRpcClientConfig grpcConf; + grpcConf.Locator = IamEndpoint_.Endpoint; + grpcConf.EnableSsl = true; + Connection_ = THolder<TServiceConnection<TService>>(Client->CreateGRpcServiceConnection<TService>(grpcConf).release()); + } + + void UpdateTicket(bool sync = false) { + with_lock(Lock_) { + if (NeedStop_ || RequestInflight_) { + return; + } + RequestInflight_ = true; + } + + auto resultPromise = NThreading::NewPromise(); + + std::shared_ptr<TImpl> self = TGrpcIamCredentialsProvider<TRequest, TResponse, TService>::TImpl::shared_from_this(); + + auto cb = [self, resultPromise, sync]( + NGrpc::TGrpcStatus&& status, TResponse&& result) mutable { + self->ProcessIamResponse(std::move(status), std::move(result), sync); + resultPromise.SetValue(); + }; + + TRequest req; + + RequestFiller_(req); + + Connection_->template DoRequest<TRequest, TResponse>( + std::move(req), + std::move(cb), + &TService::Stub::AsyncCreate, + { {}, {}, IamEndpoint_.RequestTimeout } + ); + + if (sync) { + resultPromise.GetFuture().Wait(2 * IamEndpoint_.RequestTimeout); + } + } + + TStringType GetTicket() { + TInstant nextTicketUpdate; + TString ticket; + with_lock(Lock_) { + ticket = Ticket_; + nextTicketUpdate = NextTicketUpdate_; + if (ticket.empty()) + ythrow yexception() << "IAM-token not ready yet. " << LastRequestError_; + } + if (TInstant::Now() >= nextTicketUpdate) { + UpdateTicket(); + } + return ticket; + } + + void Stop() { + with_lock(Lock_) { + if (NeedStop_) { + return; + } + NeedStop_ = true; + } + + Client.Reset(); // Will trigger destroy + } + + private: + void ProcessIamResponse(NGrpc::TGrpcStatus&& status, TResponse&& result, bool sync) { + if (!status.Ok()) { + TDuration sleepDuration; + with_lock(Lock_) { + LastRequestError_ = TStringBuilder() + << "Last request error was at " << TInstant::Now() + << ". GrpcStatusCode: " << status.GRpcStatusCode + << " Message: \"" << status.Msg + << "\" internal: " << status.InternalError + << " iam-endpoint: \"" << IamEndpoint_.Endpoint << "\""; + + RequestInflight_ = false; + sleepDuration = std::min(BackoffTimeout_, BACKOFF_MAX); + BackoffTimeout_ *= 2; + } + + Sleep(sleepDuration); + + UpdateTicket(sync); + } else { + with_lock(Lock_) { + LastRequestError_ = ""; + Ticket_ = result.iam_token(); + RequestInflight_ = false; + BackoffTimeout_ = BACKOFF_START; + + const auto now = Now(); + NextTicketUpdate_ = std::min( + now + IamEndpoint_.RefreshPeriod, + TInstant::Seconds(result.expires_at().seconds()) + ) - IamEndpoint_.RequestTimeout; + NextTicketUpdate_ = std::max(NextTicketUpdate_, now + TDuration::MilliSeconds(100)); + } + } + } + + private: + + THolder<TGRpcClientLow> Client; + THolder<TServiceConnection<TService>> Connection_; + TStringType Ticket_; + TInstant NextTicketUpdate_; + const TIamEndpoint IamEndpoint_; + const TRequestFiller RequestFiller_; + bool RequestInflight_; + TStringType LastRequestError_; + bool NeedStop_; + TDuration BackoffTimeout_; + TAdaptiveLock Lock_; + }; + +public: + TGrpcIamCredentialsProvider(const TIamEndpoint& endpoint, const TRequestFiller& requestFiller) + : Impl_(std::make_shared<TImpl>(endpoint, requestFiller)) + { + Impl_->UpdateTicket(true); + } + + ~TGrpcIamCredentialsProvider() { + Impl_->Stop(); + } + + TStringType GetAuthInfo() const override { + return Impl_->GetTicket(); + } + + bool IsValid() const override { + return true; + } + +private: + std::shared_ptr<TImpl> Impl_; +}; + +struct TIamJwtParams : TIamEndpoint { + TJwtParams JwtParams; +}; + +template<typename TRequest, typename TResponse, typename TService> +class TIamJwtCredentialsProvider : public TGrpcIamCredentialsProvider<TRequest, TResponse, TService> { +public: + TIamJwtCredentialsProvider(const TIamJwtParams& params) + : TGrpcIamCredentialsProvider<TRequest, TResponse, TService>(params, + [jwtParams = params.JwtParams](TRequest& req) { + req.set_jwt(MakeSignedJwt(jwtParams)); + }) {} +}; + +class TIamOAuthCredentialsProvider : public TGrpcIamCredentialsProvider<CreateIamTokenRequest, CreateIamTokenResponse, IamTokenService> { +public: + TIamOAuthCredentialsProvider(const TIamOAuth& params) + : TGrpcIamCredentialsProvider(params, + [token = params.OAuthToken](CreateIamTokenRequest& req) { + req.set_yandex_passport_oauth_token(token); + }) {} +}; + +TJwtParams ReadJwtKeyFile(const TString& filename) { + return ParseJwtParams(TFileInput(filename).ReadAll()); +} + +class TIamCredentialsProviderFactory : public ICredentialsProviderFactory { +public: + TIamCredentialsProviderFactory(const TIamHost& params): Params_(params) {} + + TCredentialsProviderPtr CreateProvider() const final { + return std::make_shared<TIAMCredentialsProvider>(Params_); + } + + TStringType GetClientIdentity() const final { + return "IAM_PROVIDER" + ToString((ui64)this); + } + +private: + TIamHost Params_; +}; + +template<typename TRequest, typename TResponse, typename TService> +class TIamJwtCredentialsProviderFactory : public ICredentialsProviderFactory { +public: + TIamJwtCredentialsProviderFactory(const TIamJwtParams& params): Params_(params) {} + + TCredentialsProviderPtr CreateProvider() const final { + return std::make_shared<TIamJwtCredentialsProvider<TRequest, TResponse, TService>>(Params_); + } + + TStringType GetClientIdentity() const final { + return "IAM_JWT_PROVIDER" + ToString((ui64)this); + } + +private: + TIamJwtParams Params_; +}; + +class TIamOAuthCredentialsProviderFactory : public ICredentialsProviderFactory { +public: + TIamOAuthCredentialsProviderFactory(const TIamOAuth& params): Params_(params) {} + + TCredentialsProviderPtr CreateProvider() const final { + return std::make_shared<TIamOAuthCredentialsProvider>(Params_); + } + + TStringType GetClientIdentity() const final { + return "IAM_OAUTH_PROVIDER" + ToString((ui64)this); + } + +private: + TIamOAuth Params_; +}; + +TCredentialsProviderFactoryPtr CreateIamCredentialsProviderFactory(const TIamHost& params) { + return std::make_shared<TIamCredentialsProviderFactory>(params); +} + +TCredentialsProviderFactoryPtr CreateIamJwtCredentialsProviderFactoryImpl(TIamJwtParams&& jwtParams, bool usePrivateApi) { + if (usePrivateApi) { + return std::make_shared<TIamJwtCredentialsProviderFactory< + yandex::cloud::priv::iam::v1::CreateIamTokenRequest, + yandex::cloud::priv::iam::v1::CreateIamTokenResponse, + yandex::cloud::priv::iam::v1::IamTokenService + >>(std::move(jwtParams)); + } + return std::make_shared<TIamJwtCredentialsProviderFactory<CreateIamTokenRequest, + CreateIamTokenResponse, + IamTokenService>>(std::move(jwtParams)); + +} + +TCredentialsProviderFactoryPtr CreateIamJwtFileCredentialsProviderFactory(const TIamJwtFilename& params, bool usePrivateApi) { + TIamJwtParams jwtParams = { params, ReadJwtKeyFile(params.JwtFilename) }; + return CreateIamJwtCredentialsProviderFactoryImpl(std::move(jwtParams), usePrivateApi); +} + +TCredentialsProviderFactoryPtr CreateIamJwtParamsCredentialsProviderFactory(const TIamJwtContent& params, bool usePrivateApi) { + TIamJwtParams jwtParams = { params, ParseJwtParams(params.JwtContent) }; + return CreateIamJwtCredentialsProviderFactoryImpl(std::move(jwtParams), usePrivateApi); +} + +TCredentialsProviderFactoryPtr CreateIamOAuthCredentialsProviderFactory(const TIamOAuth& params) { + return std::make_shared<TIamOAuthCredentialsProviderFactory>(params); +} +} // namespace NYdb diff --git a/kikimr/public/sdk/cpp/client/iam/iam.h b/kikimr/public/sdk/cpp/client/iam/iam.h new file mode 100644 index 0000000000..68d6d16db5 --- /dev/null +++ b/kikimr/public/sdk/cpp/client/iam/iam.h @@ -0,0 +1,50 @@ +#pragma once + +#include <ydb/public/sdk/cpp/client/ydb_types/credentials/credentials.h> + +#include <ydb/public/lib/jwt/jwt.h> +#include <util/datetime/base.h> + +namespace NYdb { + +namespace NIam { +constexpr TStringBuf DEFAULT_ENDPOINT = "iam.api.cloud.yandex.net"; + +constexpr TStringBuf DEFAULT_HOST = "169.254.169.254"; +constexpr ui32 DEFAULT_PORT = 80; + +constexpr TDuration DEFAULT_REFRESH_PERIOD = TDuration::Hours(1); +constexpr TDuration DEFAULT_REQUEST_TIMEOUT = TDuration::Seconds(10); +} + +struct TIamHost { + TString Host = TString(NIam::DEFAULT_HOST); + ui32 Port = NIam::DEFAULT_PORT; + TDuration RefreshPeriod = NIam::DEFAULT_REFRESH_PERIOD; +}; + +struct TIamEndpoint { + TString Endpoint = TString(NIam::DEFAULT_ENDPOINT); + TDuration RefreshPeriod = NIam::DEFAULT_REFRESH_PERIOD; + TDuration RequestTimeout = NIam::DEFAULT_REQUEST_TIMEOUT; +}; + +struct TIamJwtFilename : TIamEndpoint { TString JwtFilename; }; + +struct TIamJwtContent : TIamEndpoint { TString JwtContent; }; + +struct TIamOAuth : TIamEndpoint { TString OAuthToken; }; + +/// Acquire an IAM token using a local metadata service on a virtual machine. +TCredentialsProviderFactoryPtr CreateIamCredentialsProviderFactory(const TIamHost& params = {}); + +/// Acquire an IAM token using a JSON Web Token (JWT) file name. +TCredentialsProviderFactoryPtr CreateIamJwtFileCredentialsProviderFactory(const TIamJwtFilename& params, bool usePrivateApi = false); + +/// Acquire an IAM token using JSON Web Token (JWT) contents. +TCredentialsProviderFactoryPtr CreateIamJwtParamsCredentialsProviderFactory(const TIamJwtContent& param, bool usePrivateApi = false); + +// Acquire an IAM token using a user OAuth token. +TCredentialsProviderFactoryPtr CreateIamOAuthCredentialsProviderFactory(const TIamOAuth& params); + +} // namespace NYdb diff --git a/kikimr/public/sdk/cpp/client/ydb_persqueue/CMakeLists.txt b/kikimr/public/sdk/cpp/client/ydb_persqueue/CMakeLists.txt new file mode 100644 index 0000000000..9450a17202 --- /dev/null +++ b/kikimr/public/sdk/cpp/client/ydb_persqueue/CMakeLists.txt @@ -0,0 +1,16 @@ + +# This file was gererated by the build system used internally in the Yandex monorepo. +# Only simple modifications are allowed (adding source-files to targets, adding simple properties +# like target_include_directories). These modifications will be ported to original +# ya.make files by maintainers. Any complex modifications which can't be ported back to the +# original buildsystem will not be accepted. + + + +add_library(cpp-client-ydb_persqueue INTERFACE) +target_link_libraries(cpp-client-ydb_persqueue INTERFACE + contrib-libs-cxxsupp + yutil + client-ydb_persqueue-codecs + cpp-client-ydb_persqueue_core +) diff --git a/kikimr/public/sdk/cpp/client/ydb_persqueue/codecs/CMakeLists.txt b/kikimr/public/sdk/cpp/client/ydb_persqueue/codecs/CMakeLists.txt new file mode 100644 index 0000000000..e8ac82ab8c --- /dev/null +++ b/kikimr/public/sdk/cpp/client/ydb_persqueue/codecs/CMakeLists.txt @@ -0,0 +1,22 @@ + +# This file was gererated by the build system used internally in the Yandex monorepo. +# Only simple modifications are allowed (adding source-files to targets, adding simple properties +# like target_include_directories). These modifications will be ported to original +# ya.make files by maintainers. Any complex modifications which can't be ported back to the +# original buildsystem will not be accepted. + + + +add_library(client-ydb_persqueue-codecs) +target_link_libraries(client-ydb_persqueue-codecs PUBLIC + contrib-libs-cxxsupp + yutil + cpp-streams-lzop + cpp-streams-zstd + public-issue-protos + api-grpc-draft + api-protos +) +target_sources(client-ydb_persqueue-codecs PRIVATE + ${CMAKE_SOURCE_DIR}/kikimr/public/sdk/cpp/client/ydb_persqueue/codecs/codecs.cpp +) diff --git a/kikimr/public/sdk/cpp/client/ydb_persqueue/codecs/codecs.cpp b/kikimr/public/sdk/cpp/client/ydb_persqueue/codecs/codecs.cpp new file mode 100644 index 0000000000..6c0bdd1215 --- /dev/null +++ b/kikimr/public/sdk/cpp/client/ydb_persqueue/codecs/codecs.cpp @@ -0,0 +1,82 @@ +#include <library/cpp/streams/lzop/lzop.h> +#include <library/cpp/streams/zstd/zstd.h> +#include <util/stream/buffer.h> +#include <util/stream/zlib.h> +#include <util/stream/mem.h> + +#include "codecs.h" + +namespace NYdb::NPersQueue { +namespace NCompressionDetails { + +using TInputStreamVariant = std::variant<std::monostate, TZLibDecompress, TLzopDecompress, TZstdDecompress>; + +IInputStream* CreateDecompressorStream(TInputStreamVariant& inputStreamStorage, Ydb::PersQueue::V1::Codec codec, IInputStream* origin) { + switch (codec) { + case Ydb::PersQueue::V1::CODEC_GZIP: + return &inputStreamStorage.emplace<TZLibDecompress>(origin); + case Ydb::PersQueue::V1::CODEC_LZOP: + return &inputStreamStorage.emplace<TLzopDecompress>(origin); + case Ydb::PersQueue::V1::CODEC_ZSTD: + return &inputStreamStorage.emplace<TZstdDecompress>(origin); + default: + //case Ydb::PersQueue::V1::CODEC_RAW: + //case Ydb::PersQueue::V1::CODEC_UNSPECIFIED: + throw yexception() << "unsupported codec value : " << ui64(codec); + } +} + +TString Decompress(const Ydb::PersQueue::V1::MigrationStreamingReadServerMessage::DataBatch::MessageData& data) { + TMemoryInput input(data.data().data(), data.data().size()); + TString result; + TStringOutput resultOutput(result); + TInputStreamVariant inputStreamStorage; + TransferData(CreateDecompressorStream(inputStreamStorage, data.codec(), &input), &resultOutput); + return result; +} + + +class TZLibToStringCompressor: private TEmbedPolicy<TBufferOutput>, public TZLibCompress { +public: + TZLibToStringCompressor(TBuffer& dst, ZLib::StreamType type, size_t quality) + : TEmbedPolicy<TBufferOutput>(dst) + , TZLibCompress(TEmbedPolicy::Ptr(), type, quality) + { + } +}; + +class TLzopToStringCompressor: private TEmbedPolicy<TBufferOutput>, public TLzopCompress { +public: + TLzopToStringCompressor(TBuffer& dst) + : TEmbedPolicy<TBufferOutput>(dst) + , TLzopCompress(TEmbedPolicy::Ptr()) + { + } +}; + +class TZstdToStringCompressor: private TEmbedPolicy<TBufferOutput>, public TZstdCompress { +public: + TZstdToStringCompressor(TBuffer& dst, int quality) + : TEmbedPolicy<TBufferOutput>(dst) + , TZstdCompress(TEmbedPolicy::Ptr(), quality) + { + } +}; + +THolder<IOutputStream> CreateCoder(ECodec codec, TBuffer& result, int quality) { + switch (codec) { + case ECodec::GZIP: + return MakeHolder<TZLibToStringCompressor>(result, ZLib::GZip, quality >= 0 ? quality : 6); + case ECodec::LZOP: + return MakeHolder<TLzopToStringCompressor>(result); + case ECodec::ZSTD: + return MakeHolder<TZstdToStringCompressor>(result, quality); + default: + Y_FAIL("NOT IMPLEMENTED CODEC TYPE"); + } +} + + +} // namespace NDecompressionDetails + +} // namespace NYdb::NPersQueue diff --git a/kikimr/public/sdk/cpp/client/ydb_persqueue/codecs/codecs.h b/kikimr/public/sdk/cpp/client/ydb_persqueue/codecs/codecs.h new file mode 100644 index 0000000000..5fb5bb466f --- /dev/null +++ b/kikimr/public/sdk/cpp/client/ydb_persqueue/codecs/codecs.h @@ -0,0 +1,18 @@ +#pragma once +#include <util/stream/output.h> +#include <ydb/public/api/protos/ydb_persqueue_v1.pb.h> +#include <kikimr/public/sdk/cpp/client/ydb_persqueue/persqueue.h> + + +namespace NYdb::NPersQueue { +namespace NCompressionDetails { + + +extern TString Decompress(const Ydb::PersQueue::V1::MigrationStreamingReadServerMessage::DataBatch::MessageData& data); + +THolder<IOutputStream> CreateCoder(ECodec codec, TBuffer& result, int quality); + +} // namespace NDecompressionDetails + +} // namespace NYdb::NPersQueue + diff --git a/kikimr/public/sdk/cpp/client/ydb_persqueue/persqueue.h b/kikimr/public/sdk/cpp/client/ydb_persqueue/persqueue.h new file mode 100644 index 0000000000..44ba01d94f --- /dev/null +++ b/kikimr/public/sdk/cpp/client/ydb_persqueue/persqueue.h @@ -0,0 +1,2 @@ +#pragma once +#include <ydb/public/sdk/cpp/client/ydb_persqueue_core/persqueue.h> diff --git a/kikimr/yndx/api/grpc/CMakeLists.txt b/kikimr/yndx/api/grpc/CMakeLists.txt new file mode 100644 index 0000000000..adcf9dcf68 --- /dev/null +++ b/kikimr/yndx/api/grpc/CMakeLists.txt @@ -0,0 +1,43 @@ + +# This file was gererated by the build system used internally in the Yandex monorepo. +# Only simple modifications are allowed (adding source-files to targets, adding simple properties +# like target_include_directories). These modifications will be ported to original +# ya.make files by maintainers. Any complex modifications which can't be ported back to the +# original buildsystem will not be accepted. + + + +add_library(api-grpc-yndx) +set_property(TARGET api-grpc-yndx PROPERTY + PROTOC_EXTRA_OUTS .grpc.pb.cc .grpc.pb.h +) +target_link_libraries(api-grpc-yndx PUBLIC + contrib-libs-cxxsupp + yutil + contrib-libs-grpc + api-protos-yndx + api-protos + contrib-libs-protobuf +) +target_proto_messages(api-grpc-yndx PRIVATE + ${CMAKE_SOURCE_DIR}/kikimr/yndx/api/grpc/persqueue.proto + ${CMAKE_SOURCE_DIR}/kikimr/yndx/api/grpc/ydb_yndx_keyvalue_v1.proto + ${CMAKE_SOURCE_DIR}/kikimr/yndx/api/grpc/ydb_yndx_rate_limiter_v1.proto +) +target_proto_addincls(api-grpc-yndx + ./ + ${CMAKE_SOURCE_DIR}/ + ${CMAKE_BINARY_DIR} + ${CMAKE_SOURCE_DIR} + ${CMAKE_SOURCE_DIR}/contrib/libs/protobuf/src + ${CMAKE_BINARY_DIR} + ${CMAKE_SOURCE_DIR}/contrib/libs/protobuf/src +) +target_proto_outs(api-grpc-yndx + --cpp_out=${CMAKE_BINARY_DIR}/ + --cpp_styleguide_out=${CMAKE_BINARY_DIR}/ +) +target_proto_plugin(api-grpc-yndx + grpc_cpp + grpc_cpp +) diff --git a/kikimr/yndx/api/grpc/persqueue.proto b/kikimr/yndx/api/grpc/persqueue.proto new file mode 100644 index 0000000000..313b42ca3d --- /dev/null +++ b/kikimr/yndx/api/grpc/persqueue.proto @@ -0,0 +1,68 @@ +syntax = "proto3"; + +package NPersQueue; + +option java_package = "com.yandex.persqueue"; +option java_outer_classname = "PersqueueGrpc"; + +import "kikimr/yndx/api/protos/persqueue.proto"; + +service PersQueueService { + + /** + * Creates Write Session + * Pipeline: + * client server + * Init(Topic, SourceId, ...) + * ----------------> + * Init(Partition, MaxSeqNo, ...) + * <---------------- + * write(data1, seqNo1) + * ----------------> + * write(data2, seqNo2) + * ----------------> + * ack(seqNo1, offset1, ...) + * <---------------- + * write(data3, seqNo3) + * ----------------> + * ack(seqNo2, offset2, ...) + * <---------------- + * error(description, errorCode) + * <---------------- + */ + + rpc WriteSession(stream WriteRequest) returns (stream WriteResponse); + + /** + * Creates Read Session + * Pipeline: + * client server + * Init(Topics, ClientId, ...) + * ----------------> + * Init(SessionId) + * <---------------- + * read1 + * ----------------> + * read2 + * ----------------> + * lock(Topic1,Partition1, ...) - locks and releases are optional + * <---------------- + * lock(Topic2, Partition2, ...) + * <---------------- + * release(Topic1, Partition1, ...) + * <---------------- + * locked(Topic2, Partition2, ...) - client must respond to lock request with this message. Only after this client will start recieving messages from this partition + * ----------------> + * read result(data, ...) + * <---------------- + * commit(cookie1) + * ----------------> + * commit result(cookie1) + * <---------------- + * error(description, errorCode) + * <---------------- + */ + + rpc ReadSession(stream ReadRequest) returns (stream ReadResponse); + +}
\ No newline at end of file diff --git a/kikimr/yndx/api/grpc/ydb_yndx_keyvalue_v1.proto b/kikimr/yndx/api/grpc/ydb_yndx_keyvalue_v1.proto new file mode 100644 index 0000000000..7b685f5966 --- /dev/null +++ b/kikimr/yndx/api/grpc/ydb_yndx_keyvalue_v1.proto @@ -0,0 +1,43 @@ +syntax = "proto3"; + +package Ydb.Yndx.KeyValue.V1; + +option java_package = "com.yandex.ydb.yndx.keyvalue.v1"; +option java_outer_classname = "KeyValueGrpc"; +option java_multiple_files = true; + +import "kikimr/yndx/api/protos/ydb_yndx_keyvalue.proto"; + +// KeyValue tablets provide a simple key-value storage in a low-overhead and easy-to-shoot-your-leg manner. +// To use KeyValue tablets in an efficient way one must be familiar with the design of both the KeyValue tablet +// and the Distributed Storage underneath it. + +service KeyValueService { + + // Create a keyvalue volume by the path and a count of partitions + rpc CreateVolume(KeyValue.CreateVolumeRequest) returns (KeyValue.CreateVolumeResponse); + + // Drop the keyvalue volume by the path + rpc DropVolume(KeyValue.DropVolumeRequest) returns (KeyValue.DropVolumeResponse); + + // List partitions of keyvalue volume in the local node. + rpc ListLocalPartitions(KeyValue.ListLocalPartitionsRequest) returns (KeyValue.ListLocalPartitionsResponse); + + // Obtains an exclusive lock for the tablet. + rpc AcquireLock(KeyValue.AcquireLockRequest) returns (KeyValue.AcquireLockResponse); + + // Performs one or more actions that modify the state of the tablet as an atomic transaction. + rpc ExecuteTransaction(KeyValue.ExecuteTransactionRequest) returns (KeyValue.ExecuteTransactionResponse); + + // Reads value stored in the item with the key specified. + rpc Read(KeyValue.ReadRequest) returns (KeyValue.ReadResponse); + + // Reads a list of items with the keys in the range specified. + rpc ReadRange(KeyValue.ReadRangeRequest) returns (KeyValue.ReadRangeResponse); + + // List existed items with the keys in the range specified. + rpc ListRange(KeyValue.ListRangeRequest) returns (KeyValue.ListRangeResponse); + + // Gets storage channel status of the tablet. + rpc GetStorageChannelStatus(KeyValue.GetStorageChannelStatusRequest) returns (KeyValue.GetStorageChannelStatusResponse); +} diff --git a/kikimr/yndx/api/grpc/ydb_yndx_rate_limiter_v1.proto b/kikimr/yndx/api/grpc/ydb_yndx_rate_limiter_v1.proto new file mode 100644 index 0000000000..9e0c2e97dd --- /dev/null +++ b/kikimr/yndx/api/grpc/ydb_yndx_rate_limiter_v1.proto @@ -0,0 +1,35 @@ +syntax = "proto3"; + +package Ydb.Yndx.RateLimiter.V1; + +option java_package = "com.yandex.ydb.yndx.rate_limiter.v1"; +option java_outer_classname = "YndxRateLimiterGrpc"; +option java_multiple_files = true; + +import "kikimr/yndx/api/protos/ydb_yndx_rate_limiter.proto"; + +// Service that implements distributed rate limiting and accounting. +// +// To use rate limiter functionality you need an existing coordination node. + +service YndxRateLimiterService { + // Control plane API + + // Create a new resource in existing coordination node. + rpc CreateResource(CreateResourceRequest) returns (CreateResourceResponse); + + // Update a resource in coordination node. + rpc AlterResource(AlterResourceRequest) returns (AlterResourceResponse); + + // Delete a resource from coordination node. + rpc DropResource(DropResourceRequest) returns (DropResourceResponse); + + // List resources in given coordination node. + rpc ListResources(ListResourcesRequest) returns (ListResourcesResponse); + + // Describe properties of resource in coordination node. + rpc DescribeResource(DescribeResourceRequest) returns (DescribeResourceResponse); + + // Take units for usage of a resource in coordination node. + rpc AcquireResource(AcquireResourceRequest) returns (AcquireResourceResponse); +} diff --git a/kikimr/yndx/api/protos/CMakeLists.txt b/kikimr/yndx/api/protos/CMakeLists.txt new file mode 100644 index 0000000000..3f858c162c --- /dev/null +++ b/kikimr/yndx/api/protos/CMakeLists.txt @@ -0,0 +1,40 @@ + +# This file was gererated by the build system used internally in the Yandex monorepo. +# Only simple modifications are allowed (adding source-files to targets, adding simple properties +# like target_include_directories). These modifications will be ported to original +# ya.make files by maintainers. Any complex modifications which can't be ported back to the +# original buildsystem will not be accepted. + + + +add_library(api-protos-yndx) +target_link_libraries(api-protos-yndx PUBLIC + contrib-libs-cxxsupp + yutil + api-protos + tools-enum_parser-enum_serialization_runtime + contrib-libs-protobuf +) +target_proto_messages(api-protos-yndx PRIVATE + ${CMAKE_SOURCE_DIR}/kikimr/yndx/api/protos/persqueue.proto + ${CMAKE_SOURCE_DIR}/kikimr/yndx/api/protos/ydb_yndx_keyvalue.proto + ${CMAKE_SOURCE_DIR}/kikimr/yndx/api/protos/ydb_yndx_rate_limiter.proto +) +generate_enum_serilization(api-protos-yndx + ${CMAKE_BINARY_DIR}/kikimr/yndx/api/protos/persqueue.pb.h + INCLUDE_HEADERS + kikimr/yndx/api/protos/persqueue.pb.h +) +target_proto_addincls(api-protos-yndx + ./ + ${CMAKE_SOURCE_DIR}/ + ${CMAKE_BINARY_DIR} + ${CMAKE_SOURCE_DIR} + ${CMAKE_SOURCE_DIR}/contrib/libs/protobuf/src + ${CMAKE_BINARY_DIR} + ${CMAKE_SOURCE_DIR}/contrib/libs/protobuf/src +) +target_proto_outs(api-protos-yndx + --cpp_out=${CMAKE_BINARY_DIR}/ + --cpp_styleguide_out=${CMAKE_BINARY_DIR}/ +) diff --git a/kikimr/yndx/api/protos/persqueue.proto b/kikimr/yndx/api/protos/persqueue.proto new file mode 100644 index 0000000000..e8da4c4a36 --- /dev/null +++ b/kikimr/yndx/api/protos/persqueue.proto @@ -0,0 +1,335 @@ +syntax = "proto3"; +import "google/protobuf/descriptor.proto"; +import "ydb/public/api/protos/draft/persqueue_common.proto"; + +package NPersQueue; + +option java_package = "com.yandex.ydb.persqueue"; +option cc_enable_arenas = true; + +extend google.protobuf.FileOptions { + bool GenerateYaStyle = 66677; +} + +message Path { + // Path of object (topic/consumer). + string path = 1; +} + +// WRITE REQUEST + +message KeyValue { + string key = 1; + string value = 2; +} + +message MapType { + repeated KeyValue items = 1; +} + +/** + * Request for write session. Contains one of : + * Init - consists of initialization info - Topic, SourceId and so on + * Data - data to be writen + * DataBatch - batch of data to be written + */ +message WriteRequest { + message Init { + string topic = 1; + bytes source_id = 2; + + MapType extra_fields = 7; //server and file inside here + + uint64 proxy_cookie = 8; //cookie provided by ChooseProxy request //change to bytes + + uint32 partition_group = 12; //Group to write to - 0 means any; + + string version = 999; //must be filled by client lib + } + + message Data { + uint64 seq_no = 1; + bytes data = 2; + uint64 create_time_ms = 3; //timestamp in ms + NPersQueueCommon.ECodec codec = 4; + uint32 uncompressed_size = 5; + } + + message DataBatch { + repeated Data data = 1; + } + + oneof request { + //init must be sent as first message + Init init = 1; + Data data = 2; + DataBatch data_batch = 3; + } + + NPersQueueCommon.Credentials credentials = 20; +} + +/** + * Response for write session. Contains one of : + * Error - in any error state - grpc errors, session dies, incorrect Init request and so on + * Init - contains SessionId of created session, MaxSeqNo and Partition + * Ack - acknowlegment of storing corresponding message + * AckBatch - acknowlegment of storing corresponding message batch + */ +message WriteResponse { + message Init { + uint64 max_seq_no = 1; + string session_id = 2; + uint32 partition = 3; + string topic = 4; + } + + message Stat { + uint32 write_time_ms = 1; + uint32 total_time_in_partition_queue_ms = 2; + uint32 partition_quoted_time_ms = 3; + uint32 topic_quoted_time_ms = 4; + } + + message Ack { + uint64 seq_no = 1; + uint64 offset = 2; + bool already_written = 3; + + Stat stat = 4; //not filled in batch case + } + + message AckBatch { + Stat stat = 2; //common statistics for batch storing + + repeated Ack ack = 1; + } + + oneof response { + Init init = 1; + Ack ack = 2; + AckBatch ack_batch = 4; + NPersQueueCommon.Error error = 3; + } +} + +// READ REQUEST + +/** + * Request for read session. Contains one of : + * Init - contains of Topics to be readed, ClientId and other metadata + * Read - request for read batch. Contains of restrictments for result - MaxSize, MaxCount and so on + * Commit - request for commit some read batches. Contains corresponding cookies + * Locked - comfirming to server that client is ready to get data from partition from concreet offset + */ +message ReadRequest { + enum EProtocolVersion { + Base = 0; // Base protocol version + Batching = 1; // Client supports more effective batching structs (TBatchedData instead of TData) + ReadParamsInInit = 2; // Client sets read params in Init request + } + + message Init { + repeated string topics = 1; + bool read_only_local = 2; // ReadOnlyLocal=false - read mirrored topics from other clusters too; will be renamed to read_only_original + + string client_id = 4; + bool clientside_locks_allowed = 5; //if true then partitions Lock signal will be sent from server, + //and reads from partitions will began only after Locked signal recieved by server from client + + uint64 proxy_cookie = 6; //cookie provided by ChooseProxy request + + bool balance_partition_right_now = 8; //if set then do not wait for commits from client on data from partition in case of balancing + + repeated uint32 partition_groups = 9; //Groups to be read - if empty then read from all of them + + uint32 idle_timeout_sec = 10; //TODO: do we need it? + + uint32 commit_interval_ms = 12; // How often server must commit data. If client sends commits faster, + // then server will hold them in order to archive corresponding rate; zero means server default = 1sec + + // Read request params + uint32 max_read_messages_count = 14; // Max messages to give to client in one read request + uint32 max_read_size = 15; // Max size in bytes to give to client in one read request + uint32 max_read_partitions_count = 16; // 0 means not matters // Maximum partitions count to give to client in one read request + uint32 max_time_lag_ms = 17; // Read data only with time lag less than or equal to specified + uint64 read_timestamp_ms = 18; // Read data only after this timestamp + + bool commits_disabled = 19; // Client will never commit + + string version = 999; //must be filled by client lib + + // Protocol version to let server know about new features that client supports + uint32 protocol_version = 13; // version must be integer (not enum) because client may be newer than server + } + + message Read { + // It is not allowed to change these parameters. + // They will be removed in future from TRead structure. + uint32 max_count = 1; + uint32 max_size = 2; + uint32 partitions_at_once = 3; //0 means not matters + uint32 max_time_lag_ms = 5; + uint64 read_timestamp_ms = 6; //read data only after this timestamp + } + + message StartRead { + string topic = 1; + uint32 partition = 2; + + uint64 read_offset = 3; //skip upto this position; if committed position is bigger, then do nothing + bool verify_read_offset = 4; //if true then check that committed position is <= ReadOffset; otherwise it means error in client logic + uint64 generation = 5; + uint64 commit_offset = 6; //all messages BEFORE this position are processed by client + } + + message Commit { + repeated uint64 cookie = 1; + } + + message Status { + uint64 generation = 1; + string topic = 2; + uint32 partition = 3; + } + + oneof request { + //init must be sent as first message + Init init = 1; + Read read = 2; + StartRead start_read = 3; + Commit commit = 4; + Status status = 5; + } + + NPersQueueCommon.Credentials credentials = 20; +} + + +message MessageMeta { + bytes source_id = 1; + uint64 seq_no = 2; + uint64 create_time_ms = 3; + uint64 write_time_ms = 4; + + MapType extra_fields = 7; + NPersQueueCommon.ECodec codec = 8; + string ip = 9; + uint32 uncompressed_size = 10; +} + +/** + * Response for read session. Contains one of : + * Error - in any error state - grpc errors, session dies, incorrect Init request and so on + * Init - contains SessionId of created session + * Data - result of read, contains of messages batch and cookie + * Commit - acknowlegment for commit + * Lock - informs client that server is ready to read data from corresponding partition + * Release - informs client that server will not get data from this partition in future read results, unless other Lock-Locked conversation will be done + */ + +message ReadResponse { + message Init { + string session_id = 2; //for debug only + } + + message Data { + message Message { + MessageMeta meta = 1; //SeqNo ... + bytes data = 2; + //unique value for clientside deduplication - Topic:Partition:Offset + uint64 offset = 3; + bytes broken_packed_data = 4; // TODO: move to pqlib + } + + message MessageBatch { + string topic = 1; + uint32 partition = 2; + repeated Message message = 3; + } + + repeated MessageBatch message_batch = 1; + uint64 cookie = 2; //Cookie to be committed by server + } + + message BatchedData { + message MessageData { + NPersQueueCommon.ECodec codec = 2; + + uint64 offset = 3; //unique value for clientside deduplication - Topic:Partition:Offset + uint64 seq_no = 4; + + uint64 create_time_ms = 5; + uint64 uncompressed_size = 6; + + bytes data = 1; + } + + message Batch { + bytes source_id = 2; + MapType extra_fields = 3; + uint64 write_time_ms = 4; + string ip = 5; + + repeated MessageData message_data = 1; + } + + message PartitionData { + string topic = 2; + uint32 partition = 3; + + repeated Batch batch = 1; + } + + uint64 cookie = 2; //Cookie to be committed by server + + repeated PartitionData partition_data = 1; //not greater than one PartitionData for each partition + } + + message Lock { + string topic = 1; + uint32 partition = 2; + + uint64 read_offset = 3; //offset to read from + uint64 end_offset = 4; //know till this time end offset + uint64 generation = 5; + } + + message Release { + string topic = 1; + uint32 partition = 2; + bool can_commit = 3; //if CanCommit=false then you can not store progress of processing data for that partition at server; + //all commits will have no effect for this partition + //if you rely on committing offsets then just drop all data for this partition without processing - another session will get them later + //if CanCommit=true and you are relying on committing offsets - you can process all data for this partition you got, + //commit cookies and be sure that no other session will ever get this data + uint64 generation = 4; + } + + message Commit { + repeated uint64 cookie = 1; //for debug purposes only + } + + // Response for status requst. + message PartitionStatus { + uint64 generation = 1; + string topic = 2; + uint32 partition = 3; + + uint64 committed_offset = 4; + uint64 end_offset = 5; + uint64 write_watermark_ms = 6; + } + + oneof response { + Init init = 1; + Data data = 2; + BatchedData batched_data = 7; + NPersQueueCommon.Error error = 3; + Lock lock = 4; + Release release = 5; + Commit commit = 6; + PartitionStatus partition_status = 8; + } +} + diff --git a/kikimr/yndx/api/protos/ydb_yndx_keyvalue.proto b/kikimr/yndx/api/protos/ydb_yndx_keyvalue.proto new file mode 100644 index 0000000000..6782a1bfe5 --- /dev/null +++ b/kikimr/yndx/api/protos/ydb_yndx_keyvalue.proto @@ -0,0 +1,460 @@ +syntax = "proto3"; +option cc_enable_arenas = true; + +package Ydb.Yndx.KeyValue; + +option java_package = "com.yandex.ydb.yndx.rate_limiter"; +option java_outer_classname = "YndxkeyValueProtos"; +option java_multiple_files = true; + +import "ydb/public/api/protos/ydb_operation.proto"; + +// +// KeyValue API. +// + +message Flags { + // Free disk space is low. + bool disk_space_cyan = 1; + + // Free disk space is low, it is recommended to stop writing additional data. + bool disk_space_light_yellow_move = 2; + bool disk_space_yellow_stop = 3; + + // Free disk space is very low, clients must stop writing additional data. + bool disk_space_light_orange = 4; + bool disk_space_orange = 5; + + // Free disk space is extremely low, operations other than deletion may not be performed. + bool disk_space_red = 6; + + // No free disk space available. + bool disk_space_black = 7; +} + + +message Statuses { + enum Status { + STATUS_UNSPECIFIED = 0; + STATUS_SUCCESS = 1; + STATUS_NO_DATA = 2; + STATUS_ERROR = 3; + STATUS_OVERRUN = 4; + } +} + + +message StorageChannel { + // XXX + Statuses.Status status = 1; + + // Storage channel index. + uint32 storage_channel = 2; + + // If present, contains the status flags of the storage channel. Empty if status flags could not be obtained. + optional Flags status_flags = 3; +} + + +message Priorities { + enum Priority { + PRIORITY_UNSPECIFIED = 0; + + // High priority for user-initiated operations. + PRIORITY_REALTIME = 1; + + // Low prioroty for background system activity. + PRIORITY_BACKGROUND = 2; + } +} + + +message KVRange { + // The first bound of the range of the keys + // If no one is assigned then specify in order for the range to begin from the lowest key + oneof from_bound { + // Specify in order for the range to include the key specified + string from_key_inclusive = 1; + // Specify in order for the range not to include the key specified + string from_key_exclusive = 2; + } + + // The second bound of the range of the keys + // If no one is assigned then specify in order for the range to end to the highest keys + oneof to_bound { + // Specify in order for the range to include the key specified + string to_key_inclusive = 3; + // Specify in order for the range not to include the key specified + string to_key_exclusive = 4; + } +} + + +message AcquireLockRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path to the volume + string path = 2; + uint64 partition_id = 3; +} + + +message AcquireLockResponse { + // Operation contains the result of the request. Check the ydb_operation.proto. + Ydb.Operations.Operation operation = 1; +} + +message AcquireLockResult { + // The generation of the lock to provide as an argument to all the operaions the user performs with the tablet. + uint64 lock_generation = 1; +} + +message ExecuteTransactionRequest { + message Command { + message Rename { + // The key to change. + string old_key = 1; + + // The new key to change the old key to. + string new_key = 2; + } + message Concat { + // Keys to use as the source for the concatenation. + repeated string input_keys = 1; + + // New key to use for the result of the concatenation + string output_key = 2; + + // Input keys are deleted after the concatenation by default. In order to keep both the inputs and the + // output, set keep_inputs to true. + bool keep_inputs = 3; + } + + // Makes a copy of a range of key-value pairs. New keys are formed by removing a prefix and/or prepending a new + // prefix. For example, copy of the key-value pairs [{aaabc,1},{aaaef,2}] can be stripped of the 'aa' prefix and + // prepended with the 'x' so that the new pairs are [{xabc, 1}, {xaef, 2}]. + message CopyRange { + // The range of keys to copy + KVRange range = 1; + + // For each source key that begins with the prefix_to_remove, that prefix is removed from the new key before + // prepending it with the prefix_to_add. Acts as filter if not empty. + string prefix_to_remove = 2; + + // The prefix_to_add prefix is prepended to each new key. + string prefix_to_add = 3; + } + message Write { + enum Tactic { + TACTIC_UNSPECIFIED = 0; + + // Write minimum required redundant data. Does not affect storage durability. + TACTIC_MAX_THROUGHPUT = 1; + + // Write additional redundant data to more disks to reduce operation duration. Does not affect storage + // durability, but will use additional space. + TACTIC_MIN_LATENCY = 2; + } + // Key of the key-value pair to write. + string key = 1; + + // Value of the key-value pair to write. + bytes value = 2; + + // Storage channel to write the value to. Channel numbers begin with 1 and may go up to approximately 250 + // (depends on the channel configuration of each tablet). + // Channel 1 is called the INLINE channel (value is stored in the index table). + // Channel 2 is called the MAIN channel (value is stored as a separate blob in the Distributed Storage). + // Channels 1 and 2 are available for all tablets. + // If the storage channel specified is not configured for the tablet, the value is stored in + // channel 2 (the MAIN channel). + uint32 storage_channel = 3; // (default = 0 is same as 2 or MAIN) + + // Priority to use for the Distributed Storage Get operation. Has no effect for the 1st (inline) storage + // channel. Defaults to PRIORITY_UNSPECIFIED which interpreted like PRIORITY_REALTIME. + Priorities.Priority priority = 4; + + // Tactic to use for the Distributed Storage Put operation. Has no effect for the 1st (inline) storage + // channel. Defaults to TACTIC_UNSPECIFIED which interpreted like TACTIC_MAX_THROUGHPUT. + Tactic tactic = 5; + } + message DeleteRange { + // The range of keys to delete + KVRange range = 1; + } + + oneof action { + // Deletes key-value pairs with keys in the range specified. + DeleteRange delete_range = 1; + + // Changes the key of a key-value pair. + Rename rename = 2; + + // Creates a copy of key-value pairs with keys in the range specified by removin and/or prepending a prefix + // specified to each key. + CopyRange copy_range = 3; + + // Creates a new key-value pair with key specified by concatenating values of multiple other key-value pairs + // with keys specified. + Concat concat = 4; + + // Creates a new key-value pair with key and value specified. + Write write = 5; + } + } + + Ydb.Operations.OperationParams operation_params = 1; + + // Path to the volume + string path = 2; + uint64 partition_id = 3; + + // Generation of the exclusive lock obtained for the tablet as a result of an AcquireLock call. + uint64 lock_generation = 4; + + // Commands to execute as a single atomic transaction. The order of execution of commands is the same as the order + // of commands in the ExecuteTransactionRequest. Order of execution of different transactions is not specified. + repeated Command commands = 5; +} + +message ExecuteTransactionResponse { + // Operation contains the result of the request. Check the ydb_operation.proto. + Ydb.Operations.Operation operation = 1; +} + +message ExecuteTransactionResult { + repeated StorageChannel storage_channel = 1; +} + +message ReadRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path to the volume + string path = 2; + uint64 partition_id = 3; + + // Generation of the exclusive lock obtained for the tablet as a result of an AcquireLock call. + uint64 lock_generation = 4; + + // Key of the key-value pair to read. + string key = 5; + + // Offset in bytes from the beginning of the value to read data from. + uint64 offset = 6; + + // Size of the data to read in bytes. 0 means "read to the end of the value". + uint64 size = 7; + + // Result protobuf size limit. If not 0, overrides the default one only with a smaller value. + uint64 limit_bytes = 8; + + // Priority to use for the Distributed Storage Get operation. Has no effect for the 1st (inline) storage + // channel. Defaults to PRIORITY_UNSPECIFIED which interpreted like PRIORITY_REALTIME. + Priorities.Priority priority = 9; +} + +message ReadResponse { + // Operation contains the result of the request. Check the ydb_operation.proto. + Ydb.Operations.Operation operation = 1; +} + +message ReadResult { + // The key of the requested key-value pair + string requested_key = 1; + + // Offset in bytes from the beginning of the value requested + uint64 requested_offset = 2; + + // Size of the data requested + uint64 requested_size = 3; + + // The bytes of the requested part of the value of the requested key-value pair + bytes value = 4; + + // XXX + string msg = 5; + + Statuses.Status status = 6; +} + +message ReadRangeRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path to the volume + string path = 2; + uint64 partition_id = 3; + + // Generation of the exclusive lock obtained for the tablet as a result of an AcquireLock call. + uint64 lock_generation = 4; + + // The range of keys to read + KVRange range = 5; + + // Result protobuf size limit. If not 0, overrides the default one only with a smaller value. + uint64 limit_bytes = 6; + + // Priority to use for the Distributed Storage Get operation. Has no effect for the 1st (inline) storage + // channel. Defaults to PRIORITY_UNSPECIFIED which interpreted like PRIORITY_REALTIME. + Priorities.Priority priority = 7; +} + +message ReadRangeResponse { + // Operation contains the result of the request. Check the ydb_operation.proto. + Ydb.Operations.Operation operation = 1; +} + +message ReadRangeResult { + message KeyValuePair { + // The key of the key-value pair. + string key = 1; + + // The value of the key-value pair. Present only if the request was performed with include_data set to true. + bytes value = 2; + + // Full size of the value of the key-value pair. + uint32 value_size = 3; + + // Unix time of the creation of the key-value pair (in ms). + uint64 creation_unix_time = 4; + + // Contains the index of the actualy used storage channel. The actually used storage channel may differ from + // the value specified in the write request for example if there were no such storage channel at the moment + // of execution of the write command. + // For values created as a result of a concatenation or a copy of such values, the storage channel of the first + // part of the value is specified. + uint32 storage_channel = 5; // Returns the _actual_ storage channel + + Statuses.Status status = 6; + } + Statuses.Status status = 1; + + // List of the key-value pairs and metadata requested. + repeated KeyValuePair pair = 2; +} + +message ListRangeRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path to the volume + string path = 2; + uint64 partition_id = 3; + + // Generation of the exclusive lock obtained for the tablet as a result of an AcquireLock call. + uint64 lock_generation = 4; + + // The range of keys to read + KVRange range = 5; + + // Result protobuf size limit. If not 0, overrides the default one only with a smaller value. + uint64 limit_bytes = 6; + + // Priority to use for the Distributed Storage Get operation. Has no effect for the 1st (inline) storage + // channel. Defaults to PRIORITY_UNSPECIFIED which interpreted like PRIORITY_REALTIME. + Priorities.Priority priority = 7; +} + +message ListRangeResponse { + // Operation contains the result of the request. Check the ydb_operation.proto. + Ydb.Operations.Operation operation = 1; +} + +message ListRangeResult { + message KeyInfo { + // The key of the key-value pair. + string key = 1; + + // Full size of the value of the key-value pair. + uint32 value_size = 2; + + // Unix time of the creation of the key-value pair (in ms). + uint64 creation_unix_time = 3; + + // Contains the index of the actualy used storage channel. The actually used storage channel may differ from + // the value specified in the write request for example if there were no such storage channel at the moment + // of execution of the write command. + // For values created as a result of a concatenation or a copy of such values, the storage channel of the first + // part of the value is specified. + uint32 storage_channel = 4; // Returns the _actual_ storage channel + } + Statuses.Status status = 1; + + // List of the key-value pairs and metadata requested. + repeated KeyInfo key = 2; +} + +message GetStorageChannelStatusRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path to the volume + string path = 2; + uint64 partition_id = 3; + + // Generation of the exclusive lock obtained for the tablet as a result of an AcquireLock call. + uint64 lock_generation = 4; + + // Storage channel index. + repeated uint32 storage_channel = 5; +} + +message GetStorageChannelStatusResponse { + // Operation contains the result of the request. Check the ydb_operation.proto. + Ydb.Operations.Operation operation = 1; +} + +message GetStorageChannelStatusResult { + repeated StorageChannel storage_channel = 1; +} + +message CreateVolumeRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path to the volume + string path = 2; + uint32 channel_profile_id = 3; + uint32 partition_count = 4; +} + +message CreateVolumeResponse { + // Operation contains the result of the request. Check the ydb_operation.proto. + Ydb.Operations.Operation operation = 1; +} + +message CreateVolumeResult { +} + +message DropVolumeRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path to the volume + string path = 2; +} + +message DropVolumeResponse { + // Operation contains the result of the request. Check the ydb_operation.proto. + Ydb.Operations.Operation operation = 1; +} + +message DropVolumeResult { +} + +message ListLocalPartitionsRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path to the volume + string path = 2; + + // If it's zero than it used local node. + uint64 node_id = 3; +} + +message ListLocalPartitionsResponse { + // Operation contains the result of the request. Check the ydb_operation.proto. + Ydb.Operations.Operation operation = 1; +} + +message ListLocalPartitionsResult { + string requested_path = 1; + uint64 node_id = 2; + + repeated uint64 partition_ids = 3; +} diff --git a/kikimr/yndx/api/protos/ydb_yndx_rate_limiter.proto b/kikimr/yndx/api/protos/ydb_yndx_rate_limiter.proto new file mode 100644 index 0000000000..10ba9627da --- /dev/null +++ b/kikimr/yndx/api/protos/ydb_yndx_rate_limiter.proto @@ -0,0 +1,273 @@ +syntax = "proto3"; +option cc_enable_arenas = true; + +package Ydb.Yndx.RateLimiter; + +option java_package = "com.yandex.ydb.yndx.rate_limiter"; +option java_outer_classname = "YndxRateLimiterProtos"; +option java_multiple_files = true; + +import "ydb/public/api/protos/ydb_operation.proto"; + +// +// Rate Limiter control API. +// + +// +// Resource properties. +// + +message AccountingConfig { + // Account consumed resources and send billing metrics. + // Default value is false (not inherited). + bool enabled = 1; + + // Period to report consumption history from clients to kesus + // Default value is inherited from parent or equals 5000 ms for root. + uint64 report_period_ms = 2; + + // Consumption history period that is sent in one message to accounting actor. + // Default value is inherited from parent or equals 1000 ms for root. + uint64 account_period_ms = 3; + + // Time window to collect data from every client. + // Any client account message that is `collect_period` late is discarded (not accounted or billed). + // Default value is inherited from parent or equals 30 seconds for root. + uint64 collect_period_sec = 4; + + // Provisioned consumption limit in units per second. + // Effective value is limited by corresponding `max_units_per_second`. + // Default value is 0 (not inherited). + double provisioned_units_per_second = 5; + + // Provisioned allowed burst equals `provisioned_coefficient * provisioned_units_per_second` units. + // Effective value is limited by corresponding PrefetchCoefficient. + // Default value is inherited from parent or equals 60 for root. + double provisioned_coefficient = 6; + + // On-demand allowed burst equals `overshoot_coefficient * prefetch_coefficient * max_units_per_second` units. + // Should be greater or equal to 1.0 + // Default value is inherited from parent or equals 1.1 for root + double overshoot_coefficient = 7; + + // Billing metric description. + message Metric { + // Send this metric to billing. + // Default value is false (not inherited). + bool enabled = 1; + + // Billing metric period (aligned to hour boundary). + // Default value is inherited from parent or equals 60 seconds for root. + uint64 billing_period_sec = 2; + + // Billing metric JSON fields (inherited from parent if not set) + string version = 3; + string schema = 5; + string cloud_id = 6; + string folder_id = 7; + string resource_id = 8; + string source_id = 9; + } + + // Consumption within provisioned limit. + // Informative metric that should be sent to billing (not billed). + Metric provisioned = 8; + + // Consumption that exceeds provisioned limit is billed as on-demand. + Metric on_demand = 9; + + // Consumption that exceeds even on-demand limit. + // Normally it is free and should not be billed. + Metric overshoot = 10; +} + +// Settings for hierarchical deficit round robin (HDRR) algorithm. +message HierarchicalDrrSettings { + // Resource consumption speed limit. + // Value is required for root resource. + // 0 is equivalent to not set. + // Must be nonnegative. + double max_units_per_second = 1; + + // Maximum burst size of resource consumption across the whole cluster + // divided by max_units_per_second. + // Default value is 1. + // This means that maximum burst size might be equal to max_units_per_second. + // 0 is equivalent to not set. + // Must be nonnegative. + double max_burst_size_coefficient = 2; + + // Prefetch in local bucket up to prefetch_coefficient*max_units_per_second units (full size). + // Default value is inherited from parent or 0.2 for root. + // Disables prefetching if any negative value is set + // (It is useful to avoid bursts in case of large number of local buckets). + double prefetch_coefficient = 3; + + // Prefetching starts if there is less than prefetch_watermark fraction of full local bucket left. + // Default value is inherited from parent or 0.75 for root. + // Must be nonnegative and less than or equal to 1. + double prefetch_watermark = 4; +} + +// Rate limiter resource description. +message Resource { + // Resource path. Elements are separated by slash. + // The first symbol is not slash. + // The first element is root resource name. + // Resource path is the path of resource inside coordination node. + string resource_path = 1; + + oneof type { + // Settings for Hierarchical DRR algorithm. + HierarchicalDrrSettings hierarchical_drr = 2; + } + + AccountingConfig accounting_config = 3; +} + +// +// CreateResource method. +// + +message CreateResourceRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path of a coordination node. + string coordination_node_path = 2; + + // Resource properties. + Resource resource = 3; +} + +message CreateResourceResponse { + // Holds CreateResourceResult in case of successful call. + Ydb.Operations.Operation operation = 1; +} + +message CreateResourceResult { +} + +// +// AlterResource method. +// + +message AlterResourceRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path of a coordination node. + string coordination_node_path = 2; + + // New resource properties. + Resource resource = 3; +} + +message AlterResourceResponse { + // Holds AlterResourceResult in case of successful call. + Ydb.Operations.Operation operation = 1; +} + +message AlterResourceResult { +} + +// +// DropResource method. +// + +message DropResourceRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path of a coordination node. + string coordination_node_path = 2; + + // Path of resource inside a coordination node. + string resource_path = 3; +} + +message DropResourceResponse { + // Holds DropResourceResult in case of successful call. + Ydb.Operations.Operation operation = 1; +} + +message DropResourceResult { +} + +// +// ListResources method. +// + +message ListResourcesRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path of a coordination node. + string coordination_node_path = 2; + + // Path of resource inside a coordination node. + // May be empty. + // In that case all root resources will be listed. + string resource_path = 3; + + // List resources recursively. + bool recursive = 4; +} + +message ListResourcesResponse { + // Holds ListResourcesResult in case of successful call. + Ydb.Operations.Operation operation = 1; +} + +message ListResourcesResult { + repeated string resource_paths = 1; +} + +// +// DescribeResource method. +// + +message DescribeResourceRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path of a coordination node. + string coordination_node_path = 2; + + // Path of resource inside a coordination node. + string resource_path = 3; +} + +message DescribeResourceResponse { + // Holds DescribeResourceResult in case of successful call. + Ydb.Operations.Operation operation = 1; +} + +message DescribeResourceResult { + Resource resource = 1; +} + +// +// AcquireResource method. +// + +message AcquireResourceRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Path of a coordination node. + string coordination_node_path = 2; + + // Path of resource inside a coordination node. + string resource_path = 3; + + oneof units { + // Request resource's units for usage. + uint64 required = 4; + + // Actually used resource's units by client. + uint64 used = 5; + } +} + +message AcquireResourceResponse { + // Holds AcquireResourceResult in case of successful call. + Ydb.Operations.Operation operation = 1; +} + +message AcquireResourceResult { +} diff --git a/kikimr/yndx/grpc_services/persqueue/CMakeLists.txt b/kikimr/yndx/grpc_services/persqueue/CMakeLists.txt new file mode 100644 index 0000000000..5b00f89620 --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/CMakeLists.txt @@ -0,0 +1,38 @@ + +# This file was gererated by the build system used internally in the Yandex monorepo. +# Only simple modifications are allowed (adding source-files to targets, adding simple properties +# like target_include_directories). These modifications will be ported to original +# ya.make files by maintainers. Any complex modifications which can't be ported back to the +# original buildsystem will not be accepted. + + + +add_library(yndx-grpc_services-persqueue) +target_link_libraries(yndx-grpc_services-persqueue PUBLIC + contrib-libs-cxxsupp + yutil + api-grpc-yndx + api-protos-yndx + yndx-persqueue-read_batch_converter + ydb-core-base + core-client-server + ydb-core-grpc_services + core-mind-address_classification + ydb-core-persqueue + core-persqueue-events + core-persqueue-writer + ydb-core-protos + ydb-library-aclib + library-persqueue-topic_parser + services-lib-actors + services-lib-sharding + ydb-services-persqueue_v1 +) +target_sources(yndx-grpc_services-persqueue PRIVATE + ${CMAKE_SOURCE_DIR}/kikimr/yndx/grpc_services/persqueue/grpc_pq_clusters_updater_actor.cpp + ${CMAKE_SOURCE_DIR}/kikimr/yndx/grpc_services/persqueue/grpc_pq_read.cpp + ${CMAKE_SOURCE_DIR}/kikimr/yndx/grpc_services/persqueue/grpc_pq_read_actor.cpp + ${CMAKE_SOURCE_DIR}/kikimr/yndx/grpc_services/persqueue/grpc_pq_write.cpp + ${CMAKE_SOURCE_DIR}/kikimr/yndx/grpc_services/persqueue/grpc_pq_write_actor.cpp + ${CMAKE_SOURCE_DIR}/kikimr/yndx/grpc_services/persqueue/persqueue.cpp +) diff --git a/kikimr/yndx/grpc_services/persqueue/grpc_pq_actor.h b/kikimr/yndx/grpc_services/persqueue/grpc_pq_actor.h new file mode 100644 index 0000000000..155c803d0d --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/grpc_pq_actor.h @@ -0,0 +1,928 @@ +#pragma once + +#include "grpc_pq_session.h" +#include "ydb/core/client/server/msgbus_server_pq_metacache.h" +#include "ydb/core/client/server/msgbus_server_persqueue.h" + +#include <ydb/core/base/events.h> +#include <ydb/core/tx/scheme_cache/scheme_cache.h> +#include <kikimr/yndx/api/grpc/persqueue.grpc.pb.h> + +#include <ydb/core/protos/grpc_pq_old.pb.h> +#include <ydb/core/protos/pqconfig.pb.h> +#include <ydb/core/persqueue/writer/source_id_encoding.h> + +#include <library/cpp/actors/core/actor_bootstrapped.h> + +#include <library/cpp/actors/core/hfunc.h> + +#include <ydb/library/persqueue/topic_parser/topic_parser.h> +#include <ydb/core/persqueue/events/global.h> +#include <ydb/core/persqueue/writer/writer.h> +#include <ydb/core/persqueue/percentile_counter.h> +#include <ydb/core/base/appdata.h> +#include <ydb/core/base/tablet_pipe.h> +#include <ydb/core/tx/tx_proxy/proxy.h> +#include <ydb/public/lib/base/msgbus_status.h> +#include <ydb/core/kqp/kqp.h> + +#include <ydb/core/base/ticket_parser.h> +#include <ydb/services/lib/actors/type_definitions.h> +#include <ydb/services/persqueue_v1/actors/read_init_auth_actor.h> +#include <ydb/services/persqueue_v1/actors/read_session_actor.h> + + +#include <util/generic/guid.h> +#include <util/system/compiler.h> + +namespace NKikimr { +namespace NGRpcProxy { + + +static inline TVector<TEvTicketParser::TEvAuthorizeTicket::TEntry> GetTicketParserEntries(const TString& dbId, const TString& folderId) { + static const TVector<TString> permissions = {"ydb.streams.write", "ydb.databases.list", + "ydb.databases.create", "ydb.databases.connect"}; + TVector<std::pair<TString, TString>> attributes; + if (!dbId.empty()) attributes.push_back({"database_id", dbId}); + if (!folderId.empty()) attributes.push_back({"folder_id", folderId}); + if (!attributes.empty()) { + return {{permissions, attributes}}; + } + return {}; +} + + + +static inline bool InternalErrorCode(NPersQueue::NErrorCode::EErrorCode errorCode) { + switch(errorCode) { + case NPersQueue::NErrorCode::UNKNOWN_TOPIC: + case NPersQueue::NErrorCode::ERROR: + case NPersQueue::NErrorCode::INITIALIZING: + case NPersQueue::NErrorCode::OVERLOAD: + case NPersQueue::NErrorCode::WRITE_ERROR_DISK_IS_FULL: + return true; + default: + return false; + } + return false; +} + + + +Ydb::StatusIds::StatusCode ConvertPersQueueInternalCodeToStatus(const NPersQueue::NErrorCode::EErrorCode code); +void FillIssue(Ydb::Issue::IssueMessage* issue, const NPersQueue::NErrorCode::EErrorCode errorCode, const TString& errorReason); + +using IWriteSessionHandlerRef = TIntrusivePtr<ISessionHandler<NPersQueue::TWriteResponse>>; +using IReadSessionHandlerRef = TIntrusivePtr<ISessionHandler<NPersQueue::TReadResponse>>; + +const TString& LocalDCPrefix(); +const TString& MirroredDCPrefix(); + +constexpr ui64 MAGIC_COOKIE_VALUE = 123456789; + +static const TDuration CHECK_ACL_DELAY = TDuration::Minutes(5); + +struct TEvPQProxy { + enum EEv { + EvWriteInit = EventSpaceBegin(TKikimrEvents::ES_PQ_PROXY), + EvWrite, + EvDone, + EvReadInit, + EvRead, + EvCloseSession, + EvPartitionReady, + EvReadResponse, + EvCommit, + EvCommitDone, + EvLocked, + EvReleasePartition, + EvPartitionReleased, + EvLockPartition, + EvRestartPipe, + EvDieCommand, + EvPartitionStatus, + EvAuth, + EvReadSessionStatus, + EvReadSessionStatusResponse, + EvDeadlineExceeded, + EvGetStatus, + EvWriteDone, + EvEnd, + }; + + struct TEvReadSessionStatus : public TEventPB<TEvReadSessionStatus, NKikimrPQ::TReadSessionStatus, EvReadSessionStatus> { + }; + + struct TEvReadSessionStatusResponse : public TEventPB<TEvReadSessionStatusResponse, NKikimrPQ::TReadSessionStatusResponse, EvReadSessionStatusResponse> { + }; + + + + struct TEvWriteInit : public NActors::TEventLocal<TEvWriteInit, EvWriteInit> { + TEvWriteInit(const NPersQueue::TWriteRequest& req, const TString& peerName, const TString& database) + : Request(req) + , PeerName(peerName) + , Database(database) + { } + + NPersQueue::TWriteRequest Request; + TString PeerName; + TString Database; + }; + + struct TEvWrite : public NActors::TEventLocal<TEvWrite, EvWrite> { + explicit TEvWrite(const NPersQueue::TWriteRequest& req) + : Request(req) + { } + + NPersQueue::TWriteRequest Request; + }; + + struct TEvDone : public NActors::TEventLocal<TEvDone, EvDone> { + TEvDone() + { } + }; + + struct TEvWriteDone : public NActors::TEventLocal<TEvWriteDone, EvWriteDone> { + TEvWriteDone(ui64 size) + : Size(size) + { } + + ui64 Size; + }; + + struct TEvReadInit : public NActors::TEventLocal<TEvReadInit, EvReadInit> { + TEvReadInit(const NPersQueue::TReadRequest& req, const TString& peerName, const TString& database) + : Request(req) + , PeerName(peerName) + , Database(database) + { } + + NPersQueue::TReadRequest Request; + TString PeerName; + TString Database; + }; + + struct TEvRead : public NActors::TEventLocal<TEvRead, EvRead> { + explicit TEvRead(const NPersQueue::TReadRequest& req, const TString& guid = CreateGuidAsString()) + : Request(req) + , Guid(guid) + { } + + NPersQueue::TReadRequest Request; + const TString Guid; + }; + struct TEvCloseSession : public NActors::TEventLocal<TEvCloseSession, EvCloseSession> { + TEvCloseSession(const TString& reason, const NPersQueue::NErrorCode::EErrorCode errorCode) + : Reason(reason) + , ErrorCode(errorCode) + { } + + const TString Reason; + NPersQueue::NErrorCode::EErrorCode ErrorCode; + }; + + struct TEvPartitionReady : public NActors::TEventLocal<TEvPartitionReady, EvPartitionReady> { + TEvPartitionReady(const NPersQueue::TTopicConverterPtr& topic, const ui32 partition, const ui64 wTime, const ui64 sizeLag, + const ui64 readOffset, const ui64 endOffset) + : Topic(topic) + , Partition(partition) + , WTime(wTime) + , SizeLag(sizeLag) + , ReadOffset(readOffset) + , EndOffset(endOffset) + { } + + NPersQueue::TTopicConverterPtr Topic; + ui32 Partition; + ui64 WTime; + ui64 SizeLag; + ui64 ReadOffset; + ui64 EndOffset; + }; + + struct TEvReadResponse : public NActors::TEventLocal<TEvReadResponse, EvReadResponse> { + explicit TEvReadResponse( + NPersQueue::TReadResponse&& resp, + ui64 nextReadOffset, + bool fromDisk, + TDuration waitQuotaTime + ) + : Response(std::move(resp)) + , NextReadOffset(nextReadOffset) + , FromDisk(fromDisk) + , WaitQuotaTime(waitQuotaTime) + { } + + NPersQueue::TReadResponse Response; + ui64 NextReadOffset; + bool FromDisk; + TDuration WaitQuotaTime; + }; + + struct TEvCommit : public NActors::TEventLocal<TEvCommit, EvCommit> { + explicit TEvCommit(ui64 readId, ui64 offset = Max<ui64>()) + : ReadId(readId) + , Offset(offset) + { } + + ui64 ReadId; + ui64 Offset; // Actual value for requests to concreete partitions + }; + + struct TEvAuth : public NActors::TEventLocal<TEvAuth, EvAuth> { + TEvAuth(const NPersQueueCommon::TCredentials& auth) + : Auth(auth) + { } + + NPersQueueCommon::TCredentials Auth; + }; + + struct TEvLocked : public NActors::TEventLocal<TEvLocked, EvLocked> { + TEvLocked(const TString& topic, ui32 partition, ui64 readOffset, ui64 commitOffset, bool verifyReadOffset, ui64 generation) + : Topic(topic) + , Partition(partition) + , ReadOffset(readOffset) + , CommitOffset(commitOffset) + , VerifyReadOffset(verifyReadOffset) + , Generation(generation) + { } + + TString Topic; + ui32 Partition; + ui64 ReadOffset; + ui64 CommitOffset; + bool VerifyReadOffset; + ui64 Generation; + }; + + struct TEvGetStatus : public NActors::TEventLocal<TEvGetStatus, EvGetStatus> { + TEvGetStatus(const TString& topic, ui32 partition, ui64 generation) + : Topic(topic) + , Partition(partition) + , Generation(generation) + { } + + TString Topic; + ui32 Partition; + ui64 Generation; + }; + + + + struct TEvCommitDone : public NActors::TEventLocal<TEvCommitDone, EvCommitDone> { + TEvCommitDone(ui64 readId, const NPersQueue::TTopicConverterPtr& topic, const ui32 partition) + : ReadId(readId) + , Topic(topic) + , Partition(partition) + { } + + ui64 ReadId; + NPersQueue::TTopicConverterPtr Topic; + ui32 Partition; + }; + + struct TEvReleasePartition : public NActors::TEventLocal<TEvReleasePartition, EvReleasePartition> { + TEvReleasePartition() + { } + }; + + struct TEvLockPartition : public NActors::TEventLocal<TEvLockPartition, EvLockPartition> { + explicit TEvLockPartition(const ui64 readOffset, const ui64 commitOffset, bool verifyReadOffset, bool startReading) + : ReadOffset(readOffset) + , CommitOffset(commitOffset) + , VerifyReadOffset(verifyReadOffset) + , StartReading(startReading) + { } + + ui64 ReadOffset; + ui64 CommitOffset; + bool VerifyReadOffset; + bool StartReading; + }; + + + struct TEvPartitionReleased : public NActors::TEventLocal<TEvPartitionReleased, EvPartitionReleased> { + TEvPartitionReleased(const NPersQueue::TTopicConverterPtr& topic, const ui32 partition) + : Topic(topic) + , Partition(partition) + { } + + NPersQueue::TTopicConverterPtr Topic; + ui32 Partition; + }; + + + struct TEvRestartPipe : public NActors::TEventLocal<TEvRestartPipe, EvRestartPipe> { + TEvRestartPipe() + { } + }; + + struct TEvDeadlineExceeded : public NActors::TEventLocal<TEvDeadlineExceeded, EvDeadlineExceeded> { + TEvDeadlineExceeded(ui64 cookie) + : Cookie(cookie) + { } + + ui64 Cookie; + }; + + + struct TEvDieCommand : public NActors::TEventLocal<TEvDieCommand, EvDieCommand> { + TEvDieCommand(const TString& reason, const NPersQueue::NErrorCode::EErrorCode errorCode) + : Reason(reason) + , ErrorCode(errorCode) + { } + + TString Reason; + NPersQueue::NErrorCode::EErrorCode ErrorCode; + }; + + struct TEvPartitionStatus : public NActors::TEventLocal<TEvPartitionStatus, EvPartitionStatus> { + TEvPartitionStatus(const NPersQueue::TTopicConverterPtr& topic, const ui32 partition, const ui64 offset, + const ui64 endOffset, ui64 writeTimestampEstimateMs, bool init = true) + : Topic(topic) + , Partition(partition) + , Offset(offset) + , EndOffset(endOffset) + , WriteTimestampEstimateMs(writeTimestampEstimateMs) + , Init(init) + { } + + NPersQueue::TTopicConverterPtr Topic; + ui32 Partition; + ui64 Offset; + ui64 EndOffset; + ui64 WriteTimestampEstimateMs; + bool Init; + }; + +}; + + + +/// WRITE ACTOR +class TWriteSessionActor : public NActors::TActorBootstrapped<TWriteSessionActor> { + using TEvDescribeTopicsRequest = NMsgBusProxy::NPqMetaCacheV2::TEvPqNewMetaCache::TEvDescribeTopicsRequest; + using TEvDescribeTopicsResponse = NMsgBusProxy::NPqMetaCacheV2::TEvPqNewMetaCache::TEvDescribeTopicsResponse; + using TPQGroupInfoPtr = TIntrusiveConstPtr<NSchemeCache::TSchemeCacheNavigate::TPQGroupInfo>; +public: + TWriteSessionActor(IWriteSessionHandlerRef handler, const ui64 cookie, const NActors::TActorId& schemeCache, + TIntrusivePtr<NMonitoring::TDynamicCounters> counters, const TString& localDC, + const TMaybe<TString> clientDC); + ~TWriteSessionActor(); + + void Bootstrap(const NActors::TActorContext& ctx); + + void Die(const NActors::TActorContext& ctx) override; + + static constexpr NKikimrServices::TActivity::EType ActorActivityType() { return NKikimrServices::TActivity::FRONT_PQ_WRITE; } +private: + STFUNC(StateFunc) { + switch (ev->GetTypeRewrite()) { + CFunc(NActors::TEvents::TSystem::Wakeup, HandleWakeup) + + HFunc(TEvTicketParser::TEvAuthorizeTicketResult, Handle); + + HFunc(TEvPQProxy::TEvDieCommand, HandlePoison) + HFunc(TEvPQProxy::TEvWriteInit, Handle) + HFunc(TEvPQProxy::TEvWrite, Handle) + HFunc(TEvPQProxy::TEvDone, Handle) + HFunc(TEvPersQueue::TEvGetPartitionIdForWriteResponse, Handle) + + HFunc(TEvDescribeTopicsResponse, Handle); + + HFunc(NPQ::TEvPartitionWriter::TEvInitResult, Handle); + HFunc(NPQ::TEvPartitionWriter::TEvWriteAccepted, Handle); + HFunc(NPQ::TEvPartitionWriter::TEvWriteResponse, Handle); + HFunc(NPQ::TEvPartitionWriter::TEvDisconnected, Handle); + HFunc(TEvTabletPipe::TEvClientDestroyed, Handle); + HFunc(TEvTabletPipe::TEvClientConnected, Handle); + + HFunc(NKqp::TEvKqp::TEvQueryResponse, Handle); + HFunc(NKqp::TEvKqp::TEvProcessResponse, Handle); + + default: + break; + }; + } + + void Handle(NKqp::TEvKqp::TEvQueryResponse::TPtr &ev, const TActorContext &ctx); + void Handle(NKqp::TEvKqp::TEvProcessResponse::TPtr &ev, const TActorContext &ctx); + + TString CheckSupportedCodec(const ui32 codecId); + void CheckACL(const TActorContext& ctx); + void InitCheckACL(const TActorContext& ctx); + void Handle(TEvTicketParser::TEvAuthorizeTicketResult::TPtr& ev, const TActorContext& ctx); + void Handle(TEvPQProxy::TEvWriteInit::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvWrite::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvDone::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPersQueue::TEvGetPartitionIdForWriteResponse::TPtr& ev, const NActors::TActorContext& ctx); + + void LogSession(const TActorContext& ctx); + + void InitAfterDiscovery(const TActorContext& ctx); + void DiscoverPartition(const NActors::TActorContext& ctx); + void SendSelectPartitionRequest(ui32 hash, const TString& topic, const NActors::TActorContext& ctx); + void UpdatePartition(const NActors::TActorContext& ctx); + void RequestNextPartition(const NActors::TActorContext& ctx); + void ProceedPartition(const ui32 partition, const NActors::TActorContext& ctx); + THolder<NKqp::TEvKqp::TEvQueryRequest> MakeUpdateSourceIdMetadataRequest(ui32 hash, const TString& topic); + + + void Handle(TEvDescribeTopicsResponse::TPtr& ev, const NActors::TActorContext& ctx); + + void Handle(NPQ::TEvPartitionWriter::TEvInitResult::TPtr& ev, const TActorContext& ctx); + void Handle(NPQ::TEvPartitionWriter::TEvWriteAccepted::TPtr& ev, const TActorContext& ctx); + void Handle(NPQ::TEvPartitionWriter::TEvWriteResponse::TPtr& ev, const TActorContext& ctx); + void Handle(NPQ::TEvPartitionWriter::TEvDisconnected::TPtr& ev, const TActorContext& ctx); + void Handle(TEvTabletPipe::TEvClientConnected::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvTabletPipe::TEvClientDestroyed::TPtr& ev, const NActors::TActorContext& ctx); + + void HandlePoison(TEvPQProxy::TEvDieCommand::TPtr& ev, const NActors::TActorContext& ctx); + void HandleWakeup(const NActors::TActorContext& ctx); + + void CloseSession(const TString& errorReason, const NPersQueue::NErrorCode::EErrorCode errorCode, const NActors::TActorContext& ctx); + + void CheckFinish(const NActors::TActorContext& ctx); + + void GenerateNextWriteRequest(const NActors::TActorContext& ctx); + + void SetupCounters(); + void SetupCounters(const TString& cloudId, const TString& dbId, const TString& folderId); + + +private: + IWriteSessionHandlerRef Handler; + + enum EState { + ES_CREATED = 1, + ES_WAIT_SCHEME_1 = 2, + ES_WAIT_SCHEME_2 = 3, + ES_WAIT_TABLE_REQUEST_1 = 4, + ES_WAIT_NEXT_PARTITION = 5, + ES_WAIT_TABLE_REQUEST_2 = 6, + ES_WAIT_TABLE_REQUEST_3 = 7, + ES_WAIT_WRITER_INIT = 8, + ES_INITED = 9, + ES_DYING = 10, + }; + + EState State; + TActorId SchemeCache; + TActorId Writer; + + TString PeerName; + TString Database; + ui64 Cookie; + + ui32 Partition; + bool PartitionFound = false; + ui32 PreferedPartition; + TString SourceId; + ui32 SelectReqsInflight = 0; + ui64 MaxSrcIdAccessTime = 0; + NPQ::NSourceIdEncoding::TEncodedSourceId EncodedSourceId; + TString OwnerCookie; + TString UserAgent; + + ui32 NumReserveBytesRequests; + + struct TWriteRequestBatchInfo: public TSimpleRefCount<TWriteRequestBatchInfo> { + using TPtr = TIntrusivePtr<TWriteRequestBatchInfo>; + + // Source requests from user (grpc session object) + std::deque<THolder<TEvPQProxy::TEvWrite>> UserWriteRequests; + + // Formed write request's size + ui64 ByteSize = 0; + + // Formed write request's cookie + ui64 Cookie = 0; + }; + + // Nonprocessed source client requests + std::deque<THolder<TEvPQProxy::TEvWrite>> Writes; + + // Formed, but not sent, batch requests to partition actor + std::deque<TWriteRequestBatchInfo::TPtr> FormedWrites; + + // Requests that is already sent to partition actor + std::deque<TWriteRequestBatchInfo::TPtr> SentMessages; + + bool WritesDone; + + THashMap<ui32, ui64> PartitionToTablet; + + TIntrusivePtr<NMonitoring::TDynamicCounters> Counters; + + NKikimr::NPQ::TMultiCounter BytesInflight; + NKikimr::NPQ::TMultiCounter BytesInflightTotal; + + ui64 BytesInflight_; + ui64 BytesInflightTotal_; + + bool NextRequestInited; + + NKikimr::NPQ::TMultiCounter SessionsCreated; + NKikimr::NPQ::TMultiCounter SessionsActive; + NKikimr::NPQ::TMultiCounter SessionsWithoutAuth; + + NKikimr::NPQ::TMultiCounter Errors; + + ui64 NextRequestCookie; + + TIntrusivePtr<NACLib::TUserToken> Token; + NPersQueueCommon::TCredentials Auth; + TString AuthStr; + bool ACLCheckInProgress; + bool FirstACLCheck; + bool ForceACLCheck; + bool RequestNotChecked; + TInstant LastACLCheckTimestamp; + TInstant LogSessionDeadline; + + ui64 BalancerTabletId; + TString DatabaseId; + TString FolderId; + TActorId PipeToBalancer; + TIntrusivePtr<TSecurityObject> SecurityObject; + TPQGroupInfoPtr PQInfo; + + NKikimrPQClient::TDataChunk InitMeta; + TString LocalDC; + TString ClientDC; + TString SelectSourceIdQuery; + TString UpdateSourceIdQuery; + TInstant LastSourceIdUpdate; + + ui64 SourceIdCreateTime = 0; + ui32 SourceIdUpdatesInflight = 0; + + + TVector<NPersQueue::TPQLabelsInfo> Aggr; + NKikimr::NPQ::TMultiCounter SLITotal; + NKikimr::NPQ::TMultiCounter SLIErrors; + TInstant StartTime; + NKikimr::NPQ::TPercentileCounter InitLatency; + NKikimr::NPQ::TMultiCounter SLIBigLatency; + + THolder<NPersQueue::TTopicNamesConverterFactory> ConverterFactory; + NPersQueue::TDiscoveryConverterPtr DiscoveryConverter; + NPersQueue::TTopicConverterPtr FullConverter; + + NPersQueue::TWriteRequest::TInit InitRequest; +}; + +class TReadSessionActor : public TActorBootstrapped<TReadSessionActor> { + using TEvDescribeTopicsRequest = NMsgBusProxy::NPqMetaCacheV2::TEvPqNewMetaCache::TEvDescribeTopicsRequest; + using TEvDescribeTopicsResponse = NMsgBusProxy::NPqMetaCacheV2::TEvPqNewMetaCache::TEvDescribeTopicsResponse; +public: + TReadSessionActor(IReadSessionHandlerRef handler, const NPersQueue::TTopicsListController& topicsHandler, const ui64 cookie, + const NActors::TActorId& schemeCache, const NActors::TActorId& newSchemeCache, TIntrusivePtr<NMonitoring::TDynamicCounters> counters, + const TMaybe<TString> clientDC); + ~TReadSessionActor(); + + void Bootstrap(const NActors::TActorContext& ctx); + + void Die(const NActors::TActorContext& ctx) override; + + static constexpr NKikimrServices::TActivity::EType ActorActivityType() { return NKikimrServices::TActivity::FRONT_PQ_READ; } + + + struct TTopicCounters { + NKikimr::NPQ::TMultiCounter PartitionsLocked; + NKikimr::NPQ::TMultiCounter PartitionsReleased; + NKikimr::NPQ::TMultiCounter PartitionsToBeReleased; + NKikimr::NPQ::TMultiCounter PartitionsToBeLocked; + NKikimr::NPQ::TMultiCounter PartitionsInfly; + NKikimr::NPQ::TMultiCounter Errors; + NKikimr::NPQ::TMultiCounter Commits; + NKikimr::NPQ::TMultiCounter WaitsForData; + }; + +private: + STFUNC(StateFunc) { + switch (ev->GetTypeRewrite()) { + CFunc(NActors::TEvents::TSystem::Wakeup, HandleWakeup) + + HFunc(NKikimr::NGRpcProxy::V1::TEvPQProxy::TEvAuthResultOk, Handle); // form auth actor + + HFunc(TEvPQProxy::TEvDieCommand, HandlePoison) + HFunc(TEvPQProxy::TEvReadInit, Handle) //from gRPC + HFunc(TEvPQProxy::TEvReadSessionStatus, Handle) // from read sessions info builder proxy + HFunc(TEvPQProxy::TEvRead, Handle) //from gRPC + HFunc(TEvPQProxy::TEvDone, Handle) //from gRPC + HFunc(TEvPQProxy::TEvWriteDone, Handle) //from gRPC + HFunc(NKikimr::NGRpcProxy::V1::TEvPQProxy::TEvCloseSession, Handle) //from partitionActor + HFunc(TEvPQProxy::TEvCloseSession, Handle) //from partitionActor + + HFunc(TEvPQProxy::TEvPartitionReady, Handle) //from partitionActor + HFunc(TEvPQProxy::TEvPartitionReleased, Handle) //from partitionActor + + HFunc(TEvPQProxy::TEvReadResponse, Handle) //from partitionActor + HFunc(TEvPQProxy::TEvCommit, Handle) //from gRPC + HFunc(TEvPQProxy::TEvLocked, Handle) //from gRPC + HFunc(TEvPQProxy::TEvGetStatus, Handle) //from gRPC + HFunc(TEvPQProxy::TEvAuth, Handle) //from gRPC + + HFunc(TEvPQProxy::TEvCommitDone, Handle) //from PartitionActor + HFunc(TEvPQProxy::TEvPartitionStatus, Handle) //from partitionActor + + HFunc(TEvPersQueue::TEvLockPartition, Handle) //from Balancer + HFunc(TEvPersQueue::TEvReleasePartition, Handle) //from Balancer + HFunc(TEvPersQueue::TEvError, Handle) //from Balancer + + HFunc(TEvTabletPipe::TEvClientDestroyed, Handle); + HFunc(TEvTabletPipe::TEvClientConnected, Handle); + + HFunc(TEvDescribeTopicsResponse, HandleDescribeTopicsResponse); + HFunc(TEvTicketParser::TEvAuthorizeTicketResult, Handle); + + default: + break; + }; + } + + void Handle(TEvPQProxy::TEvReadInit::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvReadSessionStatus::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvRead::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvReadResponse::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvDone::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvWriteDone::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(NKikimr::NGRpcProxy::V1::TEvPQProxy::TEvCloseSession::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvCloseSession::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvPartitionReady::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvPartitionReleased::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvCommit::TPtr& ev, const NActors::TActorContext& ctx); + void MakeCommit(const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvLocked::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvGetStatus::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvAuth::TPtr& ev, const NActors::TActorContext& ctx); + void ProcessAuth(const NPersQueueCommon::TCredentials& auth); + void Handle(TEvPQProxy::TEvCommitDone::TPtr& ev, const NActors::TActorContext& ctx); + void AnswerForCommitsIfCan(const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvPartitionStatus::TPtr& ev, const NActors::TActorContext& ctx); + + void Handle(TEvPersQueue::TEvLockPartition::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPersQueue::TEvReleasePartition::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPersQueue::TEvError::TPtr& ev, const NActors::TActorContext& ctx); + + void Handle(TEvTabletPipe::TEvClientConnected::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvTabletPipe::TEvClientDestroyed::TPtr& ev, const NActors::TActorContext& ctx); + [[nodiscard]] bool ProcessBalancerDead(const ui64 tabletId, const NActors::TActorContext& ctx); // returns false if actor died + + void HandlePoison(TEvPQProxy::TEvDieCommand::TPtr& ev, const NActors::TActorContext& ctx); + void HandleWakeup(const NActors::TActorContext& ctx); + void Handle(NKikimr::NGRpcProxy::V1::TEvPQProxy::TEvAuthResultOk::TPtr& ev, const NActors::TActorContext& ctx); + + void CloseSession(const TString& errorReason, const NPersQueue::NErrorCode::EErrorCode errorCode, + const NActors::TActorContext& ctx); + + void Handle(TEvTicketParser::TEvAuthorizeTicketResult::TPtr& ev, const TActorContext& ctx); + void HandleDescribeTopicsResponse(TEvDescribeTopicsResponse::TPtr& ev, const TActorContext& ctx); + + void SendAuthRequest(const TActorContext& ctx); + void CreateInitAndAuthActor(const TActorContext& ctx); + + void SetupCounters(); + void SetupTopicCounters(const NPersQueue::TTopicConverterPtr& topic); + void SetupTopicCounters(const NPersQueue::TTopicConverterPtr& topic, const TString& cloudId, const TString& dbId, + const TString& folderId); + + [[nodiscard]] bool ProcessReads(const NActors::TActorContext& ctx); // returns false if actor died + struct TFormedReadResponse; + [[nodiscard]] bool ProcessAnswer(const NActors::TActorContext& ctx, TIntrusivePtr<TFormedReadResponse> formedResponse); // returns false if actor died + + void RegisterSessions(const NActors::TActorContext& ctx); + void RegisterSession(const TActorId& pipe, const TString& topic, const TActorContext& ctx); + + struct TPartitionActorInfo; + void DropPartitionIfNeeded(THashMap<std::pair<TString, ui32>, TPartitionActorInfo>::iterator it, const TActorContext& ctx); + + bool ActualPartitionActor(const TActorId& part); + [[nodiscard]] bool ProcessReleasePartition(const THashMap<std::pair<TString, ui32>, TPartitionActorInfo>::iterator& it, + bool kill, bool couldBeReads, const TActorContext& ctx); // returns false if actor died + void InformBalancerAboutRelease(const THashMap<std::pair<TString, ui32>, TPartitionActorInfo>::iterator& it, const TActorContext& ctx); + + // returns false if check failed. + bool CheckAndUpdateReadSettings(const NPersQueue::TReadRequest::TRead& readRequest); + + static ui32 NormalizeMaxReadMessagesCount(ui32 sourceValue); + static ui32 NormalizeMaxReadSize(ui32 sourceValue); + static ui32 NormalizeMaxReadPartitionsCount(ui32 sourceValue); + + static bool RemoveEmptyMessages(NPersQueue::TReadResponse::TBatchedData& data); // returns true if there are nonempty messages + +private: + IReadSessionHandlerRef Handler; + + const TInstant StartTimestamp; + + TActorId PqMetaCache; + TActorId NewSchemeCache; + + TActorId AuthInitActor; + bool AuthInflight; + + TString InternalClientId; + TString ExternalClientId; + const TString ClientDC; + TString ClientPath; + TString Session; + TString PeerName; + TString Database; + + bool ClientsideLocksAllowed; + bool BalanceRightNow; + bool CommitsDisabled; + bool BalancersInitStarted; + + bool InitDone; + + ui32 ProtocolVersion; // from NPersQueue::TReadRequest::EProtocolVersion + // Read settings. + // Can be initialized during Init request (new preferable way) + // or during read request (old way that will be removed in future). + // These settings can't be changed (in that case server closes session). + ui32 MaxReadMessagesCount; + ui32 MaxReadSize; + ui32 MaxReadPartitionsCount; + ui32 MaxTimeLagMs; + ui64 ReadTimestampMs; + bool ReadSettingsInited; + + NPersQueueCommon::TCredentials Auth; + TString AuthStr; + TIntrusivePtr<NACLib::TUserToken> Token; + bool ForceACLCheck; + bool RequestNotChecked; + TInstant LastACLCheckTimestamp; + + struct TPartitionActorInfo { + TActorId Actor; + std::deque<ui64> Commits; + bool Reading; + bool Releasing; + bool Released; + ui64 LockGeneration; + bool LockSent; + NPersQueue::TTopicConverterPtr Converter; + + TPartitionActorInfo(const TActorId& actor, ui64 generation, const NPersQueue::TTopicConverterPtr& topic) + : Actor(actor) + , Reading(false) + , Releasing(false) + , Released(false) + , LockGeneration(generation) + , LockSent(false) + , Converter(topic) + {} + }; + + + THashSet<TActorId> ActualPartitionActors; + THashMap<std::pair<TString, ui32>, TPartitionActorInfo> Partitions; //topic[ClientSideName!]:partition -> info + + THashMap<TString, NPersQueue::TTopicConverterPtr> FullPathToConverter; // PrimaryFullPath -> Converter, for balancer replies matching + THashMap<TString, TTopicHolder> Topics; // PrimaryName ->topic info + + TVector<ui32> Groups; + bool ReadOnlyLocal; + + struct TPartitionInfo { + NPersQueue::TTopicConverterPtr Topic; + ui32 Partition; + ui64 WTime; + ui64 SizeLag; + ui64 MsgLag; + TActorId Actor; + bool operator < (const TPartitionInfo& rhs) const { + return std::tie(WTime, Topic, Partition, Actor) < std::tie(rhs.WTime, rhs.Topic, rhs.Partition, rhs.Actor); + } + }; + + TSet<TPartitionInfo> AvailablePartitions; + + struct TOffsetsInfo { + struct TPartitionOffsetInfo { + TPartitionOffsetInfo(const TActorId& sender, const TString& topic, ui32 partition, ui64 offset) + : Sender(sender) + , Topic(topic) + , Partition(partition) + , Offset(offset) + { + } + + TActorId Sender; + TString Topic; + ui32 Partition; + ui64 Offset; + }; + + // find by read id + bool operator<(ui64 readId) const { + return ReadId < readId; + } + + friend bool operator<(ui64 readId, const TOffsetsInfo& info) { + return readId < info.ReadId; + } + + ui64 ReadId = 0; + std::vector<TPartitionOffsetInfo> PartitionOffsets; + }; + + std::deque<TOffsetsInfo> Offsets; // Sequential read id -> offsets + + struct TFormedReadResponse: public TSimpleRefCount<TFormedReadResponse> { + using TPtr = TIntrusivePtr<TFormedReadResponse>; + + TFormedReadResponse(const TString& guid, const TInstant start) + : Guid(guid) + , Start(start) + , FromDisk(false) + { + } + + NPersQueue::TReadResponse Response; + ui32 RequestsInfly = 0; + i64 ByteSize = 0; + + ui64 RequestedBytes = 0; + + //returns byteSize diff + i64 ApplyResponse(NPersQueue::TReadResponse&& resp); + + TVector<NPersQueue::TReadResponse> ControlMessages; + + THashSet<TActorId> PartitionsTookPartInRead; + TSet<TPartitionInfo> PartitionsBecameAvailable; // Partitions that became available during this read request execution. + // These partitions are bringed back to AvailablePartitions after reply to this read request. + TOffsetsInfo Offsets; // Offsets without assigned read id. + + const TString Guid; + TInstant Start; + bool FromDisk; + TDuration WaitQuotaTime; + }; + + THashMap<TActorId, TFormedReadResponse::TPtr> PartitionToReadResponse; // Partition actor -> TFormedReadResponse answer that has this partition. + // PartitionsTookPartInRead in formed read response contain this actor id. + + ui64 ReadIdToResponse; + ui64 ReadIdCommitted; + TSet<ui64> NextCommits; + TInstant LastCommitTimestamp; + TDuration CommitInterval; + ui32 CommitsInfly; + + std::deque<THolder<TEvPQProxy::TEvRead>> Reads; + + ui64 Cookie; + + struct TCommitInfo { + ui64 StartReadId; + ui32 Partitions; + TInstant StartTime; + }; + + TMap<ui64, TCommitInfo> Commits; //readid->TCommitInfo + + TIntrusivePtr<NMonitoring::TDynamicCounters> Counters; + + NMonitoring::TDynamicCounters::TCounterPtr SessionsCreated; + NMonitoring::TDynamicCounters::TCounterPtr SessionsActive; + NMonitoring::TDynamicCounters::TCounterPtr SessionsWithoutAuth; + NMonitoring::TDynamicCounters::TCounterPtr SessionsWithOldBatchingVersion; // LOGBROKER-3173 + + NMonitoring::TDynamicCounters::TCounterPtr Errors; + NMonitoring::TDynamicCounters::TCounterPtr PipeReconnects; + NMonitoring::TDynamicCounters::TCounterPtr BytesInflight; + ui64 BytesInflight_; + ui64 RequestedBytes; + ui32 ReadsInfly; + + NKikimr::NPQ::TPercentileCounter PartsPerSession; + + THashMap<TString, TTopicCounters> TopicCounters; + THashMap<TString, ui32> NumPartitionsFromTopic; + + TVector<NPersQueue::TPQLabelsInfo> Aggr; + NKikimr::NPQ::TMultiCounter SLITotal; + NKikimr::NPQ::TMultiCounter SLIErrors; + TInstant StartTime; + NKikimr::NPQ::TPercentileCounter InitLatency; + NKikimr::NPQ::TPercentileCounter CommitLatency; + NKikimr::NPQ::TMultiCounter SLIBigLatency; + + NKikimr::NPQ::TPercentileCounter ReadLatency; + NKikimr::NPQ::TPercentileCounter ReadLatencyFromDisk; + NKikimr::NPQ::TMultiCounter SLIBigReadLatency; + NKikimr::NPQ::TMultiCounter ReadsTotal; + + NPersQueue::TTopicsListController TopicsHandler; + NPersQueue::TTopicsToConverter TopicsList; +}; + +} +} diff --git a/kikimr/yndx/grpc_services/persqueue/grpc_pq_clusters_updater_actor.cpp b/kikimr/yndx/grpc_services/persqueue/grpc_pq_clusters_updater_actor.cpp new file mode 100644 index 0000000000..357c535aca --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/grpc_pq_clusters_updater_actor.cpp @@ -0,0 +1,86 @@ +#include "grpc_pq_clusters_updater_actor.h" + +#include <ydb/core/base/appdata.h> +#include <ydb/core/persqueue/pq_database.h> + +namespace NKikimr { +namespace NGRpcProxy { + +static const int CLUSTERS_UPDATER_TIMEOUT_ON_ERROR = 1; + + +TClustersUpdater::TClustersUpdater(IPQClustersUpdaterCallback* callback) + : Callback(callback) + {}; + +void TClustersUpdater::Bootstrap(const NActors::TActorContext& ctx) { + ctx.Send(ctx.SelfID, new TEvPQClustersUpdater::TEvUpdateClusters()); + ctx.Send(NNetClassifier::MakeNetClassifierID(), new NNetClassifier::TEvNetClassifier::TEvSubscribe); + + Become(&TThis::StateFunc); +} + +void TClustersUpdater::Handle(TEvPQClustersUpdater::TEvUpdateClusters::TPtr&, const TActorContext &ctx) { + auto req = MakeHolder<NKqp::TEvKqp::TEvQueryRequest>(); + req->Record.MutableRequest()->SetAction(NKikimrKqp::QUERY_ACTION_EXECUTE); + req->Record.MutableRequest()->SetType(NKikimrKqp::QUERY_TYPE_SQL_DML); + req->Record.MutableRequest()->SetKeepSession(false); + req->Record.MutableRequest()->SetQuery("--!syntax_v1\nSELECT `name`, `local`, `enabled` FROM `" + AppData(ctx)->PQConfig.GetRoot() + "/Config/V2/Cluster`;"); + req->Record.MutableRequest()->SetDatabase(NKikimr::NPQ::GetDatabaseFromConfig(AppData(ctx)->PQConfig)); + req->Record.MutableRequest()->MutableTxControl()->set_commit_tx(true); + req->Record.MutableRequest()->MutableTxControl()->mutable_begin_tx()->mutable_serializable_read_write(); + ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), req.Release()); +} + +void TClustersUpdater::Handle(NNetClassifier::TEvNetClassifier::TEvClassifierUpdate::TPtr& ev, const TActorContext&) { + + Callback->NetClassifierUpdated(ev->Get()->Classifier); +} + + + + +void TClustersUpdater::Handle(NKqp::TEvKqp::TEvQueryResponse::TPtr &ev, const TActorContext &ctx) { + auto& record = ev->Get()->Record.GetRef(); + + if (record.GetYdbStatus() == Ydb::StatusIds::SUCCESS) { + auto& t = record.GetResponse().GetResults(0).GetValue().GetStruct(0); + bool local = false; + TVector<TString> clusters; + for (size_t i = 0; i < t.ListSize(); ++i) { + TString dc = t.GetList(i).GetStruct(0).GetOptional().GetText(); + local = t.GetList(i).GetStruct(1).GetOptional().GetBool(); + clusters.push_back(dc); + if (local) { + bool enabled = t.GetList(i).GetStruct(2).GetOptional().GetBool(); + Y_VERIFY(LocalCluster.empty() || LocalCluster == dc); + bool changed = LocalCluster != dc || Enabled != enabled; + if (changed) { + LocalCluster = dc; + Enabled = enabled; + Callback->CheckClusterChange(LocalCluster, Enabled); + } + } + } + if (Clusters != clusters) { + Clusters = clusters; + Callback->CheckClustersListChange(Clusters); + } + ctx.Schedule(TDuration::Seconds(AppData(ctx)->PQConfig.GetClustersUpdateTimeoutSec()), new TEvPQClustersUpdater::TEvUpdateClusters()); + } else { + LOG_ERROR_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "can't update clusters " << record); + ctx.Schedule(TDuration::Seconds(CLUSTERS_UPDATER_TIMEOUT_ON_ERROR), new TEvPQClustersUpdater::TEvUpdateClusters()); + } +} + + +void TClustersUpdater::Handle(NKqp::TEvKqp::TEvProcessResponse::TPtr &ev, const TActorContext &ctx) { + auto& record = ev->Get()->Record; + + LOG_ERROR_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "can't update clusters " << record); + ctx.Schedule(TDuration::Seconds(CLUSTERS_UPDATER_TIMEOUT_ON_ERROR), new TEvPQClustersUpdater::TEvUpdateClusters()); +} + + +} +} diff --git a/kikimr/yndx/grpc_services/persqueue/grpc_pq_clusters_updater_actor.h b/kikimr/yndx/grpc_services/persqueue/grpc_pq_clusters_updater_actor.h new file mode 100644 index 0000000000..6f1b6ade76 --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/grpc_pq_clusters_updater_actor.h @@ -0,0 +1,77 @@ +#pragma once + +#include <library/cpp/actors/core/actor_bootstrapped.h> +#include <library/cpp/actors/core/actor.h> +#include <library/cpp/actors/core/event_local.h> +#include <library/cpp/actors/core/hfunc.h> + +#include <ydb/core/base/events.h> +#include <ydb/core/kqp/kqp.h> +#include <ydb/core/mind/address_classification/net_classifier.h> + +namespace NKikimr { +namespace NGRpcProxy { + +struct TEvPQClustersUpdater { + enum EEv { + EvUpdateClusters = EventSpaceBegin(TKikimrEvents::ES_PQ_CLUSTERS_UPDATER), + EvEnd, + }; + + struct TEvUpdateClusters : public NActors::TEventLocal<TEvUpdateClusters, EvUpdateClusters> { + TEvUpdateClusters() + {} + }; +}; + +class IPQClustersUpdaterCallback { +public: + virtual ~IPQClustersUpdaterCallback() = default; + virtual void CheckClusterChange(const TString& localCluster, const bool enabled) + { + Y_UNUSED(localCluster); + Y_UNUSED(enabled); + } + + virtual void CheckClustersListChange(const TVector<TString>& clusters) + { + Y_UNUSED(clusters); + } + + virtual void NetClassifierUpdated(NAddressClassifier::TLabeledAddressClassifier::TConstPtr classifier) { + Y_UNUSED(classifier); + } +}; + +class TClustersUpdater : public NActors::TActorBootstrapped<TClustersUpdater> { +public: + TClustersUpdater(IPQClustersUpdaterCallback* callback); + + void Bootstrap(const NActors::TActorContext& ctx); + + static constexpr NKikimrServices::TActivity::EType ActorActivityType() { return NKikimrServices::TActivity::FRONT_PQ_WRITE; } // FIXME + +private: + IPQClustersUpdaterCallback* Callback; + TString LocalCluster; + TVector<TString> Clusters; + bool Enabled = false; + + STFUNC(StateFunc) { + switch (ev->GetTypeRewrite()) { + HFunc(TEvPQClustersUpdater::TEvUpdateClusters, Handle); + HFunc(NKqp::TEvKqp::TEvQueryResponse, Handle); + HFunc(NKqp::TEvKqp::TEvProcessResponse, Handle); + HFunc(NNetClassifier::TEvNetClassifier::TEvClassifierUpdate, Handle); + } + } + + void Handle(TEvPQClustersUpdater::TEvUpdateClusters::TPtr &ev, const TActorContext &ctx); + void Handle(NKqp::TEvKqp::TEvQueryResponse::TPtr &ev, const TActorContext &ctx); + void Handle(NKqp::TEvKqp::TEvProcessResponse::TPtr &ev, const TActorContext &ctx); + void Handle(NNetClassifier::TEvNetClassifier::TEvClassifierUpdate::TPtr& ev, const TActorContext& ctx); + +}; + +} +} diff --git a/kikimr/yndx/grpc_services/persqueue/grpc_pq_read.cpp b/kikimr/yndx/grpc_services/persqueue/grpc_pq_read.cpp new file mode 100644 index 0000000000..373b74bcf2 --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/grpc_pq_read.cpp @@ -0,0 +1,268 @@ +#include "grpc_pq_read.h" +#include "grpc_pq_actor.h" +#include "grpc_pq_session.h" +#include "ydb/core/client/server/grpc_proxy_status.h" + +#include <ydb/core/grpc_services/grpc_helper.h> +#include <ydb/core/tx/scheme_board/cache.h> + +using namespace NActors; +using namespace NKikimrClient; + +using grpc::Status; + +namespace NKikimr { +namespace NGRpcProxy { + +/////////////////////////////////////////////////////////////////////////////// + +using namespace NPersQueue; + +void TPQReadService::TSession::OnCreated() { + // Start waiting for new session. + Proxy->WaitReadSession(); + if (Proxy->TooMuchSessions()) { + ReplyWithError("proxy overloaded", NPersQueue::NErrorCode::OVERLOAD); + return; + } + // Create actor for current session. + auto clusters = Proxy->GetClusters(); + auto localCluster = Proxy->GetLocalCluster(); + if (NeedDiscoverClusters && (clusters.empty() || localCluster.empty())) { + //TODO: inc sli errors counter + ReplyWithError("clusters list or local cluster is empty", NPersQueue::NErrorCode::INITIALIZING); + return; + + } + if (!TopicConverterFactory->GetLocalCluster().empty()) { + TopicConverterFactory->SetLocalCluster(localCluster); + } + auto topicsHandler = std::make_unique<NPersQueue::TTopicsListController>( + TopicConverterFactory, clusters + ); + + CreateActor(std::move(topicsHandler)); + ReadyForNextRead(); +} + +void TPQReadService::TSession::OnRead(const NPersQueue::TReadRequest& request) { + switch (request.GetRequestCase()) { + case TReadRequest::kInit: { + SendEvent(new TEvPQProxy::TEvReadInit(request, GetPeerName(), GetDatabase())); + break; + } + case TReadRequest::kRead: { + SendEvent(new TEvPQProxy::TEvRead(request)); + break; + } + case TReadRequest::kStatus: { + Y_VERIFY(ActorId); + const auto& req = request.GetStatus(); + const TString& topic = req.GetTopic(); + const ui32 partition = req.GetPartition(); + const ui64 generation = req.GetGeneration(); + SendEvent(new TEvPQProxy::TEvGetStatus(topic, partition, generation)); + ReadyForNextRead(); + break; + } + case TReadRequest::kStartRead: { + Y_VERIFY(ActorId); + const auto& req = request.GetStartRead(); + const TString& topic = req.GetTopic(); + const ui32 partition = req.GetPartition(); + const ui64 readOffset = req.GetReadOffset(); + const ui64 commitOffset = req.GetCommitOffset(); + const bool verifyReadOffset = req.GetVerifyReadOffset(); + const ui64 generation = req.GetGeneration(); + + if (request.GetCredentials().GetCredentialsCase() != NPersQueueCommon::TCredentials::CREDENTIALS_NOT_SET) { + SendEvent(new TEvPQProxy::TEvAuth(request.GetCredentials())); + } + SendEvent(new TEvPQProxy::TEvLocked(topic, partition, readOffset, commitOffset, verifyReadOffset, generation)); + ReadyForNextRead(); + break; + } + case TReadRequest::kCommit: { + Y_VERIFY(ActorId); + const auto& req = request.GetCommit(); + + if (request.GetCredentials().GetCredentialsCase() != NPersQueueCommon::TCredentials::CREDENTIALS_NOT_SET) { + SendEvent(new TEvPQProxy::TEvAuth(request.GetCredentials())); + } + + // Empty cookies list will lead to no effect. + for (ui32 i = 0; i < req.CookieSize(); ++i) { + SendEvent(new TEvPQProxy::TEvCommit(req.GetCookie(i))); + } + + ReadyForNextRead(); + break; + } + + default: { + SendEvent(new TEvPQProxy::TEvCloseSession("unsupported request", NPersQueue::NErrorCode::BAD_REQUEST)); + break; + } + } +} + +void TPQReadService::TSession::OnDone() { + SendEvent(new TEvPQProxy::TEvDone()); +} + +void TPQReadService::TSession::OnWriteDone(ui64 size) { + SendEvent(new TEvPQProxy::TEvWriteDone(size)); +} + +void TPQReadService::TSession::DestroyStream(const TString& reason, const NPersQueue::NErrorCode::EErrorCode errorCode) { + // Send poison pill to the actor(if it is alive) + SendEvent(new TEvPQProxy::TEvDieCommand("read-session " + ToString<ui64>(Cookie) + ": " + reason, errorCode)); + // Remove reference to session from "cookie -> session" map. + Proxy->ReleaseSession(Cookie); +} + +bool TPQReadService::TSession::IsShuttingDown() const { + return Proxy->IsShuttingDown(); +} + +TPQReadService::TSession::TSession(std::shared_ptr<TPQReadService> proxy, + grpc::ServerCompletionQueue* cq, ui64 cookie, const TActorId& schemeCache, const TActorId& newSchemeCache, + TIntrusivePtr<NMonitoring::TDynamicCounters> counters, bool needDiscoverClusters, + const NPersQueue::TConverterFactoryPtr& converterFactory) + : ISession(cq) + , Proxy(proxy) + , Cookie(cookie) + , ActorId() + , SchemeCache(schemeCache) + , NewSchemeCache(newSchemeCache) + , Counters(counters) + , NeedDiscoverClusters(needDiscoverClusters) + , TopicConverterFactory(converterFactory) +{ +} + +void TPQReadService::TSession::Start() { + if (!Proxy->IsShuttingDown()) { + Proxy->RequestSession(&Context, &Stream, CQ, CQ, new TRequestCreated(this)); + } +} + +void TPQReadService::TSession::SendEvent(IEventBase* ev) { + Proxy->ActorSystem->Send(ActorId, ev); +} + +void TPQReadService::TSession::CreateActor(std::unique_ptr<NPersQueue::TTopicsListController>&& topicsHandler) { + auto classifier = Proxy->GetClassifier(); + + ActorId = Proxy->ActorSystem->Register( + new TReadSessionActor(this, *topicsHandler, Cookie, SchemeCache, NewSchemeCache, Counters, + classifier ? classifier->ClassifyAddress(GetPeerName()) + : "unknown")); +} + + + +ui64 TPQReadService::TSession::GetCookie() const { + return Cookie; +} + +/////////////////////////////////////////////////////////////////////////////// + + +TPQReadService::TPQReadService(NKikimr::NGRpcService::TGRpcPersQueueService* service, grpc::ServerCompletionQueue* cq, + NActors::TActorSystem* as, const TActorId& schemeCache, + TIntrusivePtr<NMonitoring::TDynamicCounters> counters, + const ui32 maxSessions) + : Service(service) + , CQ(cq) + , ActorSystem(as) + , SchemeCache(schemeCache) + , Counters(counters) + , MaxSessions(maxSessions) +{ + auto appData = ActorSystem->AppData<TAppData>(); + auto cacheCounters = GetServiceCounters(counters, "pqproxy|schemecache"); + auto cacheConfig = MakeIntrusive<NSchemeCache::TSchemeCacheConfig>(appData, cacheCounters); + NewSchemeCache = ActorSystem->Register(CreateSchemeBoardSchemeCache(cacheConfig.Get())); + // ToDo[migration]: Other conditions; + NeedDiscoverClusters = !ActorSystem->AppData<TAppData>()->PQConfig.GetTopicsAreFirstClassCitizen(); + TopicConverterFactory = std::make_shared<NPersQueue::TTopicNamesConverterFactory>( + ActorSystem->AppData<TAppData>()->PQConfig, "" + ); + + if (NeedDiscoverClusters) { + ActorSystem->Register(new TClustersUpdater(this)); + } +} + + +ui64 TPQReadService::NextCookie() { + return AtomicIncrement(LastCookie); +} + + +void TPQReadService::ReleaseSession(ui64 cookie) { + auto g(Guard(Lock)); + bool erased = Sessions.erase(cookie); + if (erased) + ActorSystem->Send(MakeGRpcProxyStatusID(ActorSystem->NodeId), new TEvGRpcProxyStatus::TEvUpdateStatus(0,0,-1,0)); + +} + +void TPQReadService::CheckClusterChange(const TString& localCluster, const bool) { + auto g(Guard(Lock)); + LocalCluster = localCluster; + TopicConverterFactory->SetLocalCluster(localCluster); +} + +void TPQReadService::NetClassifierUpdated(NAddressClassifier::TLabeledAddressClassifier::TConstPtr classifier) { + auto g(Guard(Lock)); + if (!DatacenterClassifier) { + for (auto it = Sessions.begin(); it != Sessions.end();) { + auto jt = it++; + jt->second->DestroyStream("datacenter classifier initialized, restart session please", NPersQueue::NErrorCode::INITIALIZING); + } + } + + DatacenterClassifier = classifier; +} + + +void TPQReadService::CheckClustersListChange(const TVector<TString> &clusters) { + auto g(Guard(Lock)); + Clusters = clusters; +} + +void TPQReadService::SetupIncomingRequests() { + WaitReadSession(); +} + + +void TPQReadService::WaitReadSession() { + + const ui64 cookie = NextCookie(); + + ActorSystem->Send(MakeGRpcProxyStatusID(ActorSystem->NodeId), new TEvGRpcProxyStatus::TEvUpdateStatus(0,0,1,0)); + + TSessionRef session(new TSession(shared_from_this(), CQ, cookie, SchemeCache, NewSchemeCache, Counters, + NeedDiscoverClusters, TopicConverterFactory)); + + { + auto g(Guard(Lock)); + Sessions.insert(std::make_pair(cookie, session)); + } + + session->Start(); +} + + + +bool TPQReadService::TooMuchSessions() { + auto g(Guard(Lock)); + return Sessions.size() >= MaxSessions; +} + +/////////////////////////////////////////////////////////////////////////////// + +} +} diff --git a/kikimr/yndx/grpc_services/persqueue/grpc_pq_read.h b/kikimr/yndx/grpc_services/persqueue/grpc_pq_read.h new file mode 100644 index 0000000000..9fbc177e6f --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/grpc_pq_read.h @@ -0,0 +1,146 @@ +#pragma once + +#include "persqueue.h" +#include "grpc_pq_clusters_updater_actor.h" +#include "grpc_pq_session.h" + +#include <ydb/core/client/server/grpc_base.h> +#include <ydb/library/persqueue/topic_parser/topic_parser.h> + +#include <library/cpp/grpc/server/grpc_request.h> +#include <library/cpp/actors/core/actorsystem.h> + +#include <util/generic/hash.h> +#include <util/system/mutex.h> + +namespace NKikimr { +namespace NGRpcProxy { + +class TPQReadService : public IPQClustersUpdaterCallback, public std::enable_shared_from_this<TPQReadService> { + class TSession + : public ISession<NPersQueue::TReadRequest, NPersQueue::TReadResponse> + { + + public: + void OnCreated() override; + void OnRead(const NPersQueue::TReadRequest& request) override; + void OnDone() override; + void OnWriteDone(ui64 size) override; + void DestroyStream(const TString& reason, const NPersQueue::NErrorCode::EErrorCode errorCode) override; + bool IsShuttingDown() const override; + TSession(std::shared_ptr<TPQReadService> proxy, + grpc::ServerCompletionQueue* cq, ui64 cookie, const NActors::TActorId& schemeCache, const NActors::TActorId& newSchemeCache, + TIntrusivePtr<NMonitoring::TDynamicCounters> counters, bool needDiscoverClusters, + const NPersQueue::TConverterFactoryPtr& converterFactory); + void Start() override; + void SendEvent(NActors::IEventBase* ev); + + private: + void CreateActor(std::unique_ptr<NPersQueue::TTopicsListController>&& topicsHandler); + ui64 GetCookie() const; + + private: + std::shared_ptr<TPQReadService> Proxy; + const ui64 Cookie; + + NActors::TActorId ActorId; + + const NActors::TActorId SchemeCache; + const NActors::TActorId NewSchemeCache; + + TIntrusivePtr<NMonitoring::TDynamicCounters> Counters; + bool NeedDiscoverClusters; + + NPersQueue::TConverterFactoryPtr TopicConverterFactory; + + }; + + using TSessionRef = TIntrusivePtr<TSession>; + +public: + + TPQReadService(NGRpcService::TGRpcPersQueueService* service, + grpc::ServerCompletionQueue* cq, + NActors::TActorSystem* as, const NActors::TActorId& schemeCache, TIntrusivePtr<NMonitoring::TDynamicCounters> counters, + const ui32 maxSessions); + + virtual ~TPQReadService() + {} + + void RequestSession(::grpc::ServerContext* context, ::grpc::ServerAsyncReaderWriter< ::NPersQueue::TReadResponse, ::NPersQueue::TReadRequest>* stream, + ::grpc::CompletionQueue* new_call_cq, ::grpc::ServerCompletionQueue* notification_cq, void *tag) + { + Service->GetService()->RequestReadSession(context, stream, new_call_cq, notification_cq, tag); + } + + void SetupIncomingRequests(); + + void StopService() { + AtomicSet(ShuttingDown_, 1); + } + + bool IsShuttingDown() const { + return AtomicGet(ShuttingDown_); + } + + TVector<TString> GetClusters() const { + auto g(Guard(Lock)); + return Clusters; + } + TString GetLocalCluster() const { + auto g(Guard(Lock)); + return LocalCluster; + } + + NAddressClassifier::TLabeledAddressClassifier::TConstPtr GetClassifier() const { + auto g(Guard(Lock)); + return DatacenterClassifier; + } + +private: + ui64 NextCookie(); + + void CheckClustersListChange(const TVector<TString>& clusters) override; + void CheckClusterChange(const TString& localCluster, const bool enabled) override; + void NetClassifierUpdated(NAddressClassifier::TLabeledAddressClassifier::TConstPtr classifier) override; + void UpdateTopicsHandler(); + //! Unregistry session object. + void ReleaseSession(ui64 cookie); + + //! Start listening for incoming connections. + void WaitReadSession(); + + bool TooMuchSessions(); + +private: + NKikimr::NGRpcService::TGRpcPersQueueService* Service; + + grpc::ServerContext Context; + grpc::ServerCompletionQueue* CQ; + NActors::TActorSystem* ActorSystem; + NActors::TActorId SchemeCache; + NActors::TActorId NewSchemeCache; + + TAtomic LastCookie = 0; + TMutex Lock; + THashMap<ui64, TSessionRef> Sessions; + + TVector<TString> Clusters; + TString LocalCluster; + + TIntrusivePtr<NMonitoring::TDynamicCounters> Counters; + + ui32 MaxSessions; + + TAtomic ShuttingDown_ = 0; + + NAddressClassifier::TLabeledAddressClassifier::TConstPtr DatacenterClassifier; // Detects client's datacenter by IP. May be null + + bool NeedDiscoverClusters; + NPersQueue::TConverterFactoryPtr TopicConverterFactory; + std::unique_ptr<NPersQueue::TTopicsListController> TopicsHandler; +}; + + +} +} diff --git a/kikimr/yndx/grpc_services/persqueue/grpc_pq_read_actor.cpp b/kikimr/yndx/grpc_services/persqueue/grpc_pq_read_actor.cpp new file mode 100644 index 0000000000..e9e8798713 --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/grpc_pq_read_actor.cpp @@ -0,0 +1,2585 @@ +#include "grpc_pq_actor.h" + +#include <ydb/core/base/path.h> +#include <ydb/core/client/server/msgbus_server_persqueue.h> +#include <ydb/core/protos/services.pb.h> +#include <ydb/core/persqueue/percentile_counter.h> +#include <ydb/core/persqueue/pq_database.h> +#include <ydb/core/persqueue/write_meta.h> +#include <ydb/core/persqueue/writer/source_id_encoding.h> +#include <ydb/library/persqueue/topic_parser/type_definitions.h> +#include <ydb/library/persqueue/topic_parser/topic_parser.h> +#include <ydb/library/persqueue/topic_parser/counters.h> +#include <kikimr/yndx/persqueue/read_batch_converter/read_batch_converter.h> + +#include <library/cpp/actors/core/log.h> +#include <library/cpp/actors/interconnect/interconnect.h> +#include <library/cpp/protobuf/util/repeated_field_utils.h> + +#include <util/string/strip.h> +#include <util/charset/utf8.h> + +#include <algorithm> + +using namespace NActors; +using namespace NKikimrClient; + +namespace NKikimr { + +using namespace NMsgBusProxy; + +namespace NGRpcProxy { + +using namespace NPersQueue; +using namespace NSchemeCache; +#define PQ_LOG_PREFIX "session cookie " << Cookie << " client " << InternalClientId << " session " << Session + + +//11 tries = 10,23 seconds, then each try for 5 seconds , so 21 retries will take near 1 min +static const NTabletPipe::TClientRetryPolicy RetryPolicyForPipes = { + .RetryLimitCount = 21, + .MinRetryTime = TDuration::MilliSeconds(10), + .MaxRetryTime = TDuration::Seconds(5), + .BackoffMultiplier = 2, + .DoFirstRetryInstantly = true +}; + +static const ui64 MAX_INFLY_BYTES = 25_MB; +static const ui32 MAX_INFLY_READS = 10; + +static const TDuration READ_TIMEOUT_DURATION = TDuration::Seconds(1); + +static const TDuration WAIT_DATA = TDuration::Seconds(10); +static const TDuration PREWAIT_DATA = TDuration::Seconds(9); +static const TDuration WAIT_DELTA = TDuration::MilliSeconds(500); + +static const ui64 INIT_COOKIE = Max<ui64>(); //some identifier + +static const ui32 MAX_PIPE_RESTARTS = 100; //after 100 restarts without progress kill session +static const ui32 RESTART_PIPE_DELAY_MS = 100; + +static const ui64 MAX_READ_SIZE = 100 << 20; //100mb; + +static const TDuration DEFAULT_COMMIT_RATE = TDuration::Seconds(1); //1 second; +static const ui32 MAX_COMMITS_INFLY = 3; + +static const double LAG_GROW_MULTIPLIER = 1.2; //assume that 20% more data arrived to partitions + + +//TODO: add here tracking of bytes in/out + +#define LOG_PROTO(FieldName) \ + if (proto.Has##FieldName()) { \ + res << " " << Y_STRINGIZE(FieldName) << " { " << proto.Get##FieldName().ShortDebugString() << " }"; \ + } + +#define LOG_FIELD(proto, FieldName) \ + if (proto.Has##FieldName()) { \ + res << " " << Y_STRINGIZE(FieldName) << ": " << proto.Get##FieldName(); \ + } + +TString PartitionResponseToLog(const NKikimrClient::TPersQueuePartitionResponse& proto) { + if (!proto.HasCmdReadResult()) { + return proto.ShortDebugString(); + } + TStringBuilder res; + res << "{"; + + + if (proto.CmdWriteResultSize() > 0) { + res << " CmdWriteResult {"; + for (const auto& writeRes : proto.GetCmdWriteResult()) { + res << " { " << writeRes.ShortDebugString() << " }"; + } + res << " }"; + } + + LOG_PROTO(CmdGetMaxSeqNoResult); + LOG_PROTO(CmdGetClientOffsetResult); + LOG_PROTO(CmdGetOwnershipResult); + + + if (proto.HasCmdReadResult()) { + const auto& readRes = proto.GetCmdReadResult(); + res << " CmdReadResult {"; + LOG_FIELD(readRes, MaxOffset); + LOG_FIELD(readRes, BlobsFromDisk); + LOG_FIELD(readRes, BlobsFromCache); + //LOG_FIELD(readRes, ErrorCode); + LOG_FIELD(readRes, ErrorReason); + LOG_FIELD(readRes, BlobsCachedSize); + LOG_FIELD(readRes, SizeLag); + LOG_FIELD(readRes, RealReadOffset); + if (readRes.ResultSize() > 0) { + res << " Result {"; + for (const auto &tRes: readRes.GetResult()) { + res << " {"; + LOG_FIELD(tRes, Offset); + LOG_FIELD(tRes, SeqNo); + LOG_FIELD(tRes, PartNo); + LOG_FIELD(tRes, TotalParts); + LOG_FIELD(tRes, TotalSize); + LOG_FIELD(tRes, WriteTimestampMS); + LOG_FIELD(tRes, CreateTimestampMS); + LOG_FIELD(tRes, UncompressedSize); + LOG_FIELD(tRes, PartitionKey); + res << " }"; + } + res << " }"; + } + res << " }"; + } + res << " }"; + return res; +} +#undef LOG_PROTO +#undef LOG_FIELD + +class TPartitionActor : public NActors::TActorBootstrapped<TPartitionActor> { +public: + TPartitionActor(const TActorId& parentId, const TString& clientId, const ui64 cookie, const TString& session, const ui32 generation, + const ui32 step, const NPersQueue::TTopicConverterPtr& topic, const ui32 partition, const ui64 tabletID, + const TReadSessionActor::TTopicCounters& counters, const TString& clientDC); + ~TPartitionActor(); + + void Bootstrap(const NActors::TActorContext& ctx); + void Die(const NActors::TActorContext& ctx) override; + + + static constexpr NKikimrServices::TActivity::EType ActorActivityType() { return NKikimrServices::TActivity::FRONT_PQ_PARTITION; } +private: + STFUNC(StateFunc) { + switch (ev->GetTypeRewrite()) { + CFunc(NActors::TEvents::TSystem::Wakeup, HandleWakeup) + HFunc(TEvPQProxy::TEvDeadlineExceeded, Handle) + + HFunc(NActors::TEvents::TEvPoisonPill, HandlePoison) + HFunc(TEvPQProxy::TEvRead, Handle) + HFunc(TEvPQProxy::TEvCommit, Handle) + HFunc(TEvPQProxy::TEvReleasePartition, Handle) + HFunc(TEvPQProxy::TEvLockPartition, Handle) + HFunc(TEvPQProxy::TEvGetStatus, Handle) + HFunc(TEvPQProxy::TEvRestartPipe, Handle) + + HFunc(TEvTabletPipe::TEvClientDestroyed, Handle); + HFunc(TEvTabletPipe::TEvClientConnected, Handle); + HFunc(TEvPersQueue::TEvResponse, Handle); + HFunc(TEvPersQueue::TEvHasDataInfoResponse, Handle); + default: + break; + }; + } + + + void Handle(TEvPQProxy::TEvReleasePartition::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvLockPartition::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvGetStatus::TPtr& ev, const NActors::TActorContext& ctx); + + void Handle(TEvPQProxy::TEvDeadlineExceeded::TPtr& ev, const NActors::TActorContext& ctx); + + void Handle(TEvPQProxy::TEvRead::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPQProxy::TEvCommit::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(const TEvPQProxy::TEvRestartPipe::TPtr&, const NActors::TActorContext& ctx); + + void Handle(TEvTabletPipe::TEvClientConnected::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvTabletPipe::TEvClientDestroyed::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPersQueue::TEvResponse::TPtr& ev, const NActors::TActorContext& ctx); + void Handle(TEvPersQueue::TEvHasDataInfoResponse::TPtr& ev, const NActors::TActorContext& ctx); + + void HandlePoison(NActors::TEvents::TEvPoisonPill::TPtr& ev, const NActors::TActorContext& ctx); + void HandleWakeup(const NActors::TActorContext& ctx); + + void CheckRelease(const NActors::TActorContext& ctx); + void InitLockPartition(const NActors::TActorContext& ctx); + void InitStartReading(const NActors::TActorContext& ctx); + + void RestartPipe(const NActors::TActorContext& ctx, const TString& reason, const NPersQueue::NErrorCode::EErrorCode errorCode); + void WaitDataInPartition(const NActors::TActorContext& ctx); + void SendCommit(const ui64 readId, const ui64 offset, const TActorContext& ctx); + +private: + const TActorId ParentId; + const TString InternalClientId; + const TString ClientDC; + const ui64 Cookie; + const TString Session; + const ui32 Generation; + const ui32 Step; + + NPersQueue::TTopicConverterPtr Topic; + const ui32 Partition; + + const ui64 TabletID; + + ui64 ReadOffset; + ui64 ClientReadOffset; + ui64 ClientCommitOffset; + bool ClientVerifyReadOffset; + ui64 CommittedOffset; + ui64 WriteTimestampEstimateMs; + + ui64 WTime; + bool InitDone; + bool StartReading; + bool AllPrepareInited; + bool FirstInit; + TActorId PipeClient; + ui32 PipeGeneration; + bool RequestInfly; + NKikimrClient::TPersQueueRequest CurrentRequest; + + ui64 EndOffset; + ui64 SizeLag; + + TString ReadGuid; // empty if not reading + + bool NeedRelease; + bool Released; + + std::set<ui64> WaitDataInfly; + ui64 WaitDataCookie; + bool WaitForData; + + bool LockCounted; + + std::deque<std::pair<ui64, ui64>> CommitsInfly; //ReadId, Offset + + TReadSessionActor::TTopicCounters Counters; +}; + + +TReadSessionActor::TReadSessionActor( + IReadSessionHandlerRef handler, const NPersQueue::TTopicsListController& topicsHandler, const ui64 cookie, + const TActorId& pqMetaCache, const TActorId& newSchemeCache, TIntrusivePtr<NMonitoring::TDynamicCounters> counters, + const TMaybe<TString> clientDC +) + : Handler(handler) + , StartTimestamp(TInstant::Now()) + , PqMetaCache(pqMetaCache) + , NewSchemeCache(newSchemeCache) + , AuthInitActor() + , AuthInflight(false) + , ClientDC(clientDC ? *clientDC : "other") + , ClientPath() + , Session() + , ClientsideLocksAllowed(false) + , BalanceRightNow(false) + , CommitsDisabled(false) + , BalancersInitStarted(false) + , InitDone(false) + , ProtocolVersion(NPersQueue::TReadRequest::Base) + , MaxReadMessagesCount(0) + , MaxReadSize(0) + , MaxReadPartitionsCount(0) + , MaxTimeLagMs(0) + , ReadTimestampMs(0) + , ReadSettingsInited(false) + , ForceACLCheck(false) + , RequestNotChecked(true) + , LastACLCheckTimestamp(TInstant::Zero()) + , ReadOnlyLocal(false) + , ReadIdToResponse(1) + , ReadIdCommitted(0) + , LastCommitTimestamp(TInstant::Zero()) + , CommitInterval(DEFAULT_COMMIT_RATE) + , CommitsInfly(0) + , Cookie(cookie) + , Counters(counters) + , BytesInflight_(0) + , RequestedBytes(0) + , ReadsInfly(0) + , TopicsHandler(topicsHandler) +{ + Y_ASSERT(Handler); +} + + + +TReadSessionActor::~TReadSessionActor() = default; + + +void TReadSessionActor::Bootstrap(const TActorContext& ctx) { + if (!AppData(ctx)->PQConfig.GetTopicsAreFirstClassCitizen()) { + ++(*GetServiceCounters(Counters, "pqproxy|readSession")->GetCounter("SessionsCreatedTotal", true)); + } + StartTime = ctx.Now(); + Become(&TThis::StateFunc); +} + + +void TReadSessionActor::Die(const TActorContext& ctx) { + + ctx.Send(AuthInitActor, new TEvents::TEvPoisonPill()); + + for (auto& p : Partitions) { + ctx.Send(p.second.Actor, new TEvents::TEvPoisonPill()); + + if (!p.second.Released) { + auto it = TopicCounters.find(p.second.Converter->GetInternalName()); + Y_VERIFY(it != TopicCounters.end()); + it->second.PartitionsInfly.Dec(); + it->second.PartitionsReleased.Inc(); + if (p.second.Releasing) + it->second.PartitionsToBeReleased.Dec(); + } + } + + for (auto& t : Topics) { + if (t.second.PipeClient) + NTabletPipe::CloseClient(ctx, t.second.PipeClient); + } + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " is DEAD"); + + if (SessionsActive) { + --(*SessionsActive); + } + if (BytesInflight) { + (*BytesInflight) -= BytesInflight_; + } + if (SessionsActive) { //PartsPerSession is inited too + PartsPerSession.DecFor(Partitions.size(), 1); + } + if (!Handler->IsShuttingDown()) + Handler->Finish(); + TActorBootstrapped<TReadSessionActor>::Die(ctx); +} + +void TReadSessionActor::Handle(TEvPQProxy::TEvDone::TPtr&, const TActorContext& ctx) { + CloseSession(TStringBuilder() << "Reads done signal - closing everything", NPersQueue::NErrorCode::OK, ctx); +} + +void TReadSessionActor::Handle(TEvPQProxy::TEvWriteDone::TPtr& ev, const TActorContext& ctx) { + Y_VERIFY(BytesInflight_ >= ev->Get()->Size); + BytesInflight_ -= ev->Get()->Size; + if (BytesInflight) (*BytesInflight) -= ev->Get()->Size; + + const bool isAlive = ProcessReads(ctx); + Y_UNUSED(isAlive); +} + +void TReadSessionActor::Handle(TEvPQProxy::TEvCommit::TPtr& ev, const TActorContext& ctx) { + RequestNotChecked = true; + + if (CommitsDisabled) { + CloseSession(TStringBuilder() << "commits in session are disabled by client option", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + const ui64 readId = ev->Get()->ReadId; + if (readId <= ReadIdCommitted) { + CloseSession(TStringBuilder() << "commit of " << ev->Get()->ReadId << " that is already committed", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + if (readId >= ReadIdToResponse) { + CloseSession(TStringBuilder() << "commit of unknown cookie " << readId, NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + if (NextCommits.size() >= AppData(ctx)->PQConfig.GetMaxReadCookies()) { + CloseSession(TStringBuilder() << "got more than " << AppData(ctx)->PQConfig.GetMaxReadCookies() << " unordered cookies to commit " << readId, NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + bool res = NextCommits.insert(readId).second; + if (!res) { + CloseSession(TStringBuilder() << "double commit of cookie " << readId, NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " commit request from client for " << readId); + MakeCommit(ctx); +} + +void TReadSessionActor::MakeCommit(const TActorContext& ctx) { + if (CommitsDisabled) + return; + if (ctx.Now() - LastCommitTimestamp < CommitInterval) + return; + if (CommitsInfly > MAX_COMMITS_INFLY) + return; + ui64 readId = ReadIdCommitted; + auto it = NextCommits.begin(); + for (;it != NextCommits.end() && (*it) == readId + 1; ++it) { + ++readId; + } + if (readId == ReadIdCommitted) + return; + NextCommits.erase(NextCommits.begin(), it); + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " commit request from " << ReadIdCommitted + 1 << " to " << readId); + + auto& commit = Commits[readId]; + commit.StartReadId = ReadIdCommitted + 1; + commit.Partitions = 0; + commit.StartTime = ctx.Now(); + ReadIdCommitted = readId; + LastCommitTimestamp = ctx.Now(); + ++CommitsInfly; + SLITotal.Inc(); + Y_VERIFY(Commits.size() == CommitsInfly); + + // Find last offset info belonging to our read id and its ancestors. + const auto firstGreater = std::upper_bound(Offsets.begin(), Offsets.end(), readId); + THashSet<std::pair<TString, ui64>> processedPartitions; + + // Iterate from last to first offsets to find partitions' offsets. + // Offsets in queue have nondecreasing values (for each partition), + // so it it sufficient to take only the last offset for each partition. + // Note: reverse_iterator(firstGreater) points to _before_ firstGreater + + for (auto i = std::make_reverse_iterator(firstGreater), end = std::make_reverse_iterator(Offsets.begin()); i != end; ++i) { + const TOffsetsInfo& info = *i; + for (const TOffsetsInfo::TPartitionOffsetInfo& pi : info.PartitionOffsets) { + if (!ActualPartitionActor(pi.Sender)) { + continue; + } + const auto partitionKey = std::make_pair(pi.Topic, pi.Partition); + if (!processedPartitions.insert(partitionKey).second) { + continue; // already processed + } + const auto partitionIt = Partitions.find(partitionKey); + if (partitionIt != Partitions.end() && !partitionIt->second.Released) { + ctx.Send(partitionIt->second.Actor, new TEvPQProxy::TEvCommit(readId, pi.Offset)); + partitionIt->second.Commits.push_back(readId); + ++commit.Partitions; + } + } + } + Offsets.erase(Offsets.begin(), firstGreater); + + AnswerForCommitsIfCan(ctx); //Could be done if all partitions are lost because of balancer dead +} + +void TReadSessionActor::Handle(TEvPQProxy::TEvAuth::TPtr& ev, const TActorContext&) { + ProcessAuth(ev->Get()->Auth); +} + +void TReadSessionActor::Handle(TEvPQProxy::TEvGetStatus::TPtr& ev, const TActorContext& ctx) { + + if (!ClientsideLocksAllowed) { + CloseSession("Partition status available only when ClientsideLocksAllowed is true", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + auto it = Partitions.find(std::make_pair(ev->Get()->Topic, ev->Get()->Partition)); + + if (it == Partitions.end() || it->second.Releasing || it->second.LockGeneration != ev->Get()->Generation) { + //do nothing - already released partition + LOG_WARN_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " got NOTACTUAL get status request from client for " << ev->Get()->Topic + << ":" << ev->Get()->Partition << " generation " << ev->Get()->Generation); + return; + } + + //proxy request to partition - allow initing + //TODO: add here VerifyReadOffset too and check it against Committed position + ctx.Send(it->second.Actor, new TEvPQProxy::TEvGetStatus(ev->Get()->Topic, ev->Get()->Partition, ev->Get()->Generation)); +} + + +void TReadSessionActor::Handle(TEvPQProxy::TEvLocked::TPtr& ev, const TActorContext& ctx) { + + RequestNotChecked = true; + if (!ClientsideLocksAllowed) { + CloseSession("Locked requests are allowed only when ClientsideLocksAllowed is true", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + auto& topic = ev->Get()->Topic; + if (topic.empty()) { + CloseSession("empty topic in start_read request", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + + } + auto it = Partitions.find(std::make_pair(topic, ev->Get()->Partition)); + + if (it == Partitions.end() || it->second.Releasing || it->second.LockGeneration != ev->Get()->Generation) { + //do nothing - already released partition + LOG_WARN_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " got NOTACTUAL lock from client for " << topic + << ":" << ev->Get()->Partition << " at offset " << ev->Get()->ReadOffset << " generation " << ev->Get()->Generation); + + return; + } + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " got lock from client for " << ev->Get()->Topic + << ":" << ev->Get()->Partition << " at readOffset " << ev->Get()->ReadOffset << " commitOffset " << ev->Get()->CommitOffset << " generation " << ev->Get()->Generation); + + //proxy request to partition - allow initing + //TODO: add here VerifyReadOffset too and check it against Committed position + ctx.Send(it->second.Actor, new TEvPQProxy::TEvLockPartition(ev->Get()->ReadOffset, ev->Get()->CommitOffset, ev->Get()->VerifyReadOffset, true)); +} + +void TReadSessionActor::DropPartitionIfNeeded(THashMap<std::pair<TString, ui32>, TPartitionActorInfo>::iterator it, const TActorContext& ctx) { + if (it->second.Commits.empty() && it->second.Released) { + ctx.Send(it->second.Actor, new TEvents::TEvPoisonPill()); + bool res = ActualPartitionActors.erase(it->second.Actor); + Y_VERIFY(res); + + if (--NumPartitionsFromTopic[it->second.Converter->GetInternalName()] == 0) { + bool res = TopicCounters.erase(it->second.Converter->GetInternalName()); + Y_VERIFY(res); + } + + if (SessionsActive) { + PartsPerSession.DecFor(Partitions.size(), 1); + } + Partitions.erase(it); + if (SessionsActive) { + PartsPerSession.IncFor(Partitions.size(), 1); + } + } +} + + +void TReadSessionActor::Handle(TEvPQProxy::TEvCommitDone::TPtr& ev, const TActorContext& ctx) { + + Y_VERIFY(!CommitsDisabled); + + if (!ActualPartitionActor(ev->Sender)) + return; + + ui64 readId = ev->Get()->ReadId; + + auto it = Commits.find(readId); + Y_VERIFY(it != Commits.end()); + --it->second.Partitions; + + auto jt = Partitions.find(std::make_pair(ev->Get()->Topic->GetClientsideName(), ev->Get()->Partition)); + Y_VERIFY(jt != Partitions.end()); + Y_VERIFY(!jt->second.Commits.empty() && jt->second.Commits.front() == readId); + jt->second.Commits.pop_front(); + + DropPartitionIfNeeded(jt, ctx); + + AnswerForCommitsIfCan(ctx); + + MakeCommit(ctx); +} + +void TReadSessionActor::AnswerForCommitsIfCan(const TActorContext& ctx) { + while (!Commits.empty() && Commits.begin()->second.Partitions == 0) { + auto it = Commits.begin(); + ui64 readId = it->first; + TReadResponse result; + for (ui64 i = it->second.StartReadId; i <= readId; ++i){ + result.MutableCommit()->AddCookie(i); + } + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " replying for commits from " << it->second.StartReadId + << " to " << readId); + ui64 diff = result.ByteSize(); + BytesInflight_ += diff; + if (BytesInflight) (*BytesInflight) += diff; + Handler->Reply(result); + + ui32 commitDurationMs = (ctx.Now() - it->second.StartTime).MilliSeconds(); + CommitLatency.IncFor(commitDurationMs, 1); + if (commitDurationMs >= AppData(ctx)->PQConfig.GetCommitLatencyBigMs()) { + SLIBigLatency.Inc(); + } + Commits.erase(it); + --CommitsInfly; + Y_VERIFY(Commits.size() == CommitsInfly); + } +} + + +void TReadSessionActor::Handle(TEvPQProxy::TEvReadSessionStatus::TPtr& ev, const TActorContext& ctx) { + + THolder<TEvPQProxy::TEvReadSessionStatusResponse> result(new TEvPQProxy::TEvReadSessionStatusResponse()); + result->Record.SetSession(Session); + result->Record.SetTimestamp(StartTimestamp.MilliSeconds()); + + result->Record.SetClientNode(PeerName); + result->Record.SetProxyNodeId(ctx.SelfID.NodeId()); + + for (auto& p : Partitions) { + auto part = result->Record.AddPartition(); + part->SetTopic(p.first.first); + part->SetPartition(p.first.second); + part->SetAssignId(0); + for (auto& c : NextCommits) { + part->AddNextCommits(c); + } + part->SetReadIdCommitted(ReadIdCommitted); + part->SetLastReadId(ReadIdToResponse - 1); + part->SetTimestampMs(0); + } + + ctx.Send(ev->Sender, result.Release()); +} + +void TReadSessionActor::Handle(TEvPQProxy::TEvReadInit::TPtr& ev, const TActorContext& ctx) { + + THolder<TEvPQProxy::TEvReadInit> event(ev->Release()); + + if (!Topics.empty()) { + //answer error + CloseSession("got second init request", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + const auto& init = event->Request.GetInit(); + + if (!init.TopicsSize()) { + CloseSession("no topics in init request", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + if (init.GetClientId().empty()) { + CloseSession("no clientId in init request", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + if (init.GetProxyCookie() != ctx.SelfID.NodeId() && init.GetProxyCookie() != MAGIC_COOKIE_VALUE) { + CloseSession("you must perform ChooseProxy request at first and go to ProxyName server with ProxyCookie", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + // ToDo[migration] - consider separate consumer conversion logic - ? + if (AppData(ctx)->PQConfig.GetTopicsAreFirstClassCitizen()) { + ClientPath = init.GetClientId(); + ExternalClientId = ClientPath; + InternalClientId = ConvertNewConsumerName(init.GetClientId()); + } else { + ClientPath = NormalizeFullPath(MakeConsumerPath(init.GetClientId())); + ExternalClientId = ClientPath; + InternalClientId = ConvertNewConsumerName(init.GetClientId()); + } + + Auth = event->Request.GetCredentials(); + event->Request.ClearCredentials(); + Y_PROTOBUF_SUPPRESS_NODISCARD Auth.SerializeToString(&AuthStr); + TStringBuilder session; + session << ExternalClientId << "_" << ctx.SelfID.NodeId() << "_" << Cookie << "_" << TAppData::RandomProvider->GenRand64(); + Session = session; + ProtocolVersion = init.GetProtocolVersion(); + CommitsDisabled = init.GetCommitsDisabled(); + + if (ProtocolVersion >= NPersQueue::TReadRequest::ReadParamsInInit) { + ReadSettingsInited = true; + MaxReadMessagesCount = NormalizeMaxReadMessagesCount(init.GetMaxReadMessagesCount()); + MaxReadSize = NormalizeMaxReadSize(init.GetMaxReadSize()); + MaxReadPartitionsCount = NormalizeMaxReadPartitionsCount(init.GetMaxReadPartitionsCount()); + MaxTimeLagMs = init.GetMaxTimeLagMs(); + ReadTimestampMs = init.GetReadTimestampMs(); + } + + PeerName = event->PeerName; + Database = event->Database; + + ReadOnlyLocal = init.GetReadOnlyLocal(); + + if (init.GetCommitIntervalMs()) { + CommitInterval = Min(CommitInterval, TDuration::MilliSeconds(init.GetCommitIntervalMs())); + } + + for (ui32 i = 0; i < init.PartitionGroupsSize(); ++i) { + Groups.push_back(init.GetPartitionGroups(i)); + } + THashSet<TString> topicsToResolve; + for (ui32 i = 0; i < init.TopicsSize(); ++i) { + const auto& t = init.GetTopics(i); + + if (t.empty()) { + CloseSession("empty topic in init request", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + topicsToResolve.insert(t); + } + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " init: " << event->Request << " from " << PeerName); + + ClientsideLocksAllowed = init.GetClientsideLocksAllowed(); + BalanceRightNow = init.GetBalancePartitionRightNow() || CommitsDisabled; + + if (!AppData(ctx)->PQConfig.GetTopicsAreFirstClassCitizen()) { + SetupCounters(); + } + + if (Auth.GetCredentialsCase() == NPersQueueCommon::TCredentials::CREDENTIALS_NOT_SET) { + LOG_WARN_S(ctx, NKikimrServices::PQ_READ_PROXY, "session without AuthInfo : " << ExternalClientId << " from " << PeerName); + if (SessionsWithoutAuth) { + ++(*SessionsWithoutAuth); + } + if (AppData(ctx)->PQConfig.GetRequireCredentialsInNewProtocol()) { + CloseSession("Unauthenticated access is forbidden, please provide credentials", NPersQueue::NErrorCode::ACCESS_DENIED, ctx); + return; + } + } + TopicsList = TopicsHandler.GetReadTopicsList( + topicsToResolve, ReadOnlyLocal, Database + ); + if (!TopicsList.IsValid) { + return CloseSession( + TopicsList.Reason, + NPersQueue::NErrorCode::BAD_REQUEST, ctx + ); + } + SendAuthRequest(ctx); + + auto subGroup = GetServiceCounters(Counters, "pqproxy|SLI"); + Aggr = {{{{"Account", ClientPath.substr(0, ClientPath.find("/"))}}, {"total"}}}; + SLITotal = NKikimr::NPQ::TMultiCounter(subGroup, Aggr, {}, {"RequestsTotal"}, true, "sensor", false); + SLIErrors = NKikimr::NPQ::TMultiCounter(subGroup, Aggr, {}, {"RequestsError"}, true, "sensor", false); + + SLITotal.Inc(); +} + + +void TReadSessionActor::SendAuthRequest(const TActorContext& ctx) { + AuthInitActor = {}; + AuthInflight = true; + + if (Auth.GetCredentialsCase() == NPersQueueCommon::TCredentials::CREDENTIALS_NOT_SET) { + Token = nullptr; + CreateInitAndAuthActor(ctx); + return; + } + auto database = Database.empty() ? NKikimr::NPQ::GetDatabaseFromConfig(AppData(ctx)->PQConfig) : Database; + Y_VERIFY(TopicsList.IsValid); + TVector<TDiscoveryConverterPtr> topics; + for(const auto& t : TopicsList.Topics) { + if (topics.size() >= 10) { + break; + } + topics.push_back(t.second); + } + ctx.Send(PqMetaCache, new TEvDescribeTopicsRequest(topics, false)); +} + + + +void TReadSessionActor::HandleDescribeTopicsResponse(TEvDescribeTopicsResponse::TPtr& ev, const TActorContext& ctx) { + TString dbId, folderId; + for (const auto& entry : ev->Get()->Result->ResultSet) { + Y_VERIFY(entry.PQGroupInfo); // checked at ProcessMetaCacheTopicResponse() + auto& pqDescr = entry.PQGroupInfo->Description; + dbId = pqDescr.GetPQTabletConfig().GetYdbDatabaseId(); + folderId = pqDescr.GetPQTabletConfig().GetYcFolderId(); + break; + } + + auto entries = GetTicketParserEntries(dbId, folderId); + + TString ticket; + switch (Auth.GetCredentialsCase()) { + case NPersQueueCommon::TCredentials::kTvmServiceTicket: + ticket = Auth.GetTvmServiceTicket(); + break; + case NPersQueueCommon::TCredentials::kOauthToken: + ticket = Auth.GetOauthToken(); + break; + default: + CloseSession("Unknown Credentials case", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + ctx.Send(MakeTicketParserID(), new TEvTicketParser::TEvAuthorizeTicket({ + .Database = Database, + .Ticket = ticket, + .PeerName = PeerName, + .Entries = entries + })); +} + +void TReadSessionActor::CreateInitAndAuthActor(const TActorContext& ctx) { + auto database = Database.empty() ? NKikimr::NPQ::GetDatabaseFromConfig(AppData(ctx)->PQConfig) : Database; + AuthInitActor = ctx.Register(new V1::TReadInitAndAuthActor( + ctx, ctx.SelfID, InternalClientId, Cookie, Session, PqMetaCache, NewSchemeCache, Counters, Token, + TopicsList, TopicsHandler.GetLocalCluster() + )); +} + +void TReadSessionActor::Handle(TEvTicketParser::TEvAuthorizeTicketResult::TPtr& ev, const TActorContext& ctx) { + TString ticket = ev->Get()->Ticket; + TString maskedTicket = ticket.size() > 5 ? (ticket.substr(0, 5) + "***" + ticket.substr(ticket.size() - 5)) : "***"; + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, "CheckACL ticket " << maskedTicket << " got result from TICKET_PARSER response: error: " << ev->Get()->Error << " user: " + << (ev->Get()->Error.empty() ? ev->Get()->Token->GetUserSID() : "")); + + if (!ev->Get()->Error.empty()) { + CloseSession(TStringBuilder() << "Ticket parsing error: " << ev->Get()->Error, NPersQueue::NErrorCode::ACCESS_DENIED, ctx); + return; + } + Token = ev->Get()->Token; + CreateInitAndAuthActor(ctx); +} + + +void TReadSessionActor::RegisterSession(const TActorId& pipe, const TString& topic, const TActorContext& ctx) { + + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " register session to " << topic); + THolder<TEvPersQueue::TEvRegisterReadSession> request; + request.Reset(new TEvPersQueue::TEvRegisterReadSession); + auto& req = request->Record; + req.SetSession(Session); + req.SetClientNode(PeerName); + ActorIdToProto(pipe, req.MutablePipeClient()); + req.SetClientId(InternalClientId); + + for (ui32 i = 0; i < Groups.size(); ++i) { + req.AddGroups(Groups[i]); + } + + NTabletPipe::SendData(ctx, pipe, request.Release()); +} + +void TReadSessionActor::RegisterSessions(const TActorContext& ctx) { + InitDone = true; + + for (auto& t : Topics) { + auto& topic = t.first; + RegisterSession(t.second.PipeClient, topic, ctx); + NumPartitionsFromTopic[topic] = 0; + } +} + + +void TReadSessionActor::SetupCounters() +{ + if (SessionsCreated) { + return; + } + + auto subGroup = GetServiceCounters(Counters, "pqproxy|readSession")->GetSubgroup("Client", InternalClientId)->GetSubgroup("ConsumerPath", ClientPath); + SessionsCreated = subGroup->GetExpiringCounter("SessionsCreated", true); + SessionsActive = subGroup->GetExpiringCounter("SessionsActive", false); + SessionsWithoutAuth = subGroup->GetExpiringCounter("WithoutAuth", true); + SessionsWithOldBatchingVersion = subGroup->GetExpiringCounter("SessionsWithOldBatchingVersion", true); // monitoring to ensure that old version is not used anymore + Errors = subGroup->GetExpiringCounter("Errors", true); + PipeReconnects = subGroup->GetExpiringCounter("PipeReconnects", true); + + BytesInflight = subGroup->GetExpiringCounter("BytesInflight", false); + + PartsPerSession = NKikimr::NPQ::TPercentileCounter(subGroup->GetSubgroup("sensor", "PartsPerSession"), {}, {}, "Count", + TVector<std::pair<ui64, TString>>{{1, "1"}, {2, "2"}, {5, "5"}, + {10, "10"}, {20, "20"}, {50, "50"}, {70, "70"}, + {100, "100"}, {150, "150"}, {300,"300"}, {99999999, "99999999"}}, false); + + ++(*SessionsCreated); + ++(*SessionsActive); + PartsPerSession.IncFor(Partitions.size(), 1); //for 0 + + if (ProtocolVersion < NPersQueue::TReadRequest::Batching) { + ++(*SessionsWithOldBatchingVersion); + } +} + + +void TReadSessionActor::SetupTopicCounters(const TTopicConverterPtr& topic) +{ + auto& topicCounters = TopicCounters[topic->GetInternalName()]; + auto subGroup = GetServiceCounters(Counters, "pqproxy|readSession"); +//client/consumerPath Account/Producer OriginDC Topic/TopicPath + + auto aggr = GetLabels(topic); + TVector<std::pair<TString, TString>> cons = {{"Client", InternalClientId}, {"ConsumerPath", ClientPath}}; + + topicCounters.PartitionsLocked = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"PartitionsLocked"}, true); + topicCounters.PartitionsReleased = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"PartitionsReleased"}, true); + topicCounters.PartitionsToBeReleased = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"PartitionsToBeReleased"}, false); + topicCounters.PartitionsToBeLocked = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"PartitionsToBeLocked"}, false); + topicCounters.PartitionsInfly = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"PartitionsInfly"}, false); + topicCounters.Errors = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"PartitionsErrors"}, true); + topicCounters.Commits = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"Commits"}, true); + topicCounters.WaitsForData = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"WaitsForData"}, true); +} + +void TReadSessionActor::SetupTopicCounters(const TTopicConverterPtr& topic, const TString& cloudId, + const TString& dbId, const TString& folderId) +{ + auto& topicCounters = TopicCounters[topic->GetInternalName()]; + auto subGroup = NPersQueue::GetCountersForStream(Counters); +//client/consumerPath Account/Producer OriginDC Topic/TopicPath + auto aggr = GetLabelsForStream(topic, cloudId, dbId, folderId); + TVector<std::pair<TString, TString>> cons = {{"consumer", ClientPath}}; + + topicCounters.PartitionsLocked = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"stream.internal_read.partitions_locked_per_second"}, true, "name"); + topicCounters.PartitionsReleased = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"stream.internal_read.partitions_released_per_second"}, true, "name"); + topicCounters.PartitionsToBeReleased = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"stream.internal_read.partitions_to_be_released"}, false, "name"); + topicCounters.PartitionsToBeLocked = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"stream.internal_read.partitions_to_be_locked"}, false, "name"); + topicCounters.PartitionsInfly = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"stream.internal_read.partitions_locked"}, false, "name"); + topicCounters.Errors = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"stream.internal_read.partitions_errors_per_second"}, true, "name"); + topicCounters.Commits = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"stream.internal_read.commits_per_second"}, true, "name"); + topicCounters.WaitsForData = NKikimr::NPQ::TMultiCounter(subGroup, aggr, cons, {"stream.internal_read.waits_for_data"}, true, "name"); +} + +void TReadSessionActor::Handle(V1::TEvPQProxy::TEvAuthResultOk::TPtr& ev, const TActorContext& ctx) { + + LastACLCheckTimestamp = ctx.Now(); + + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " auth ok, got " << ev->Get()->TopicAndTablets.size() << " topics, init done " << InitDone); + + AuthInitActor = TActorId(); + AuthInflight = false; + + if (!InitDone) { + + ui32 initBorder = AppData(ctx)->PQConfig.GetReadInitLatencyBigMs(); + ui32 readBorder = AppData(ctx)->PQConfig.GetReadLatencyBigMs(); + ui32 readBorderFromDisk = AppData(ctx)->PQConfig.GetReadLatencyFromDiskBigMs(); + + auto subGroup = GetServiceCounters(Counters, "pqproxy|SLI"); + InitLatency = NKikimr::NPQ::CreateSLIDurationCounter(subGroup, Aggr, "ReadInit", initBorder, {100, 200, 500, 1000, 1500, 2000, 5000, 10000, 30000, 99999999}); + CommitLatency = NKikimr::NPQ::CreateSLIDurationCounter(subGroup, Aggr, "Commit", AppData(ctx)->PQConfig.GetCommitLatencyBigMs(), {100, 200, 500, 1000, 1500, 2000, 5000, 10000, 30000, 99999999}); + SLIBigLatency = NKikimr::NPQ::TMultiCounter(subGroup, Aggr, {}, {"RequestsBigLatency"}, true, "sensor", false); + ReadLatency = NKikimr::NPQ::CreateSLIDurationCounter(subGroup, Aggr, "Read", readBorder, {100, 200, 500, 1000, 1500, 2000, 5000, 10000, 30000, 99999999}); + ReadLatencyFromDisk = NKikimr::NPQ::CreateSLIDurationCounter(subGroup, Aggr, "ReadFromDisk", readBorderFromDisk, {100, 200, 500, 1000, 1500, 2000, 5000, 10000, 30000, 99999999}); + SLIBigReadLatency = NKikimr::NPQ::TMultiCounter(subGroup, Aggr, {}, {"ReadBigLatency"}, true, "sensor", false); + ReadsTotal = NKikimr::NPQ::TMultiCounter(subGroup, Aggr, {}, {"ReadsTotal"}, true, "sensor", false); + + ui32 initDurationMs = (ctx.Now() - StartTime).MilliSeconds(); + InitLatency.IncFor(initDurationMs, 1); + if (initDurationMs >= initBorder) { + SLIBigLatency.Inc(); + } + + + TReadResponse result; + result.MutableInit()->SetSessionId(Session); + ui64 diff = result.ByteSize(); + BytesInflight_ += diff; + if (BytesInflight) (*BytesInflight) += diff; + + Handler->Reply(result); + + Handler->ReadyForNextRead(); + + Y_VERIFY(!BalancersInitStarted); + BalancersInitStarted = true; + + for (auto& t : ev->Get()->TopicAndTablets) { + auto& topicHolder = Topics[t.TopicNameConverter->GetInternalName()]; + topicHolder.TabletID = t.TabletID; + topicHolder.CloudId = t.CloudId; + topicHolder.DbId = t.DbId; + topicHolder.FolderId = t.FolderId; + topicHolder.FullConverter = t.TopicNameConverter; + FullPathToConverter[t.TopicNameConverter->GetPrimaryPath()] = t.TopicNameConverter; + const auto& second = t.TopicNameConverter->GetSecondaryPath(); + if (!second.empty()) { + FullPathToConverter[second] = t.TopicNameConverter; + } + } + + for (auto& t : Topics) { + NTabletPipe::TClientConfig clientConfig; + + clientConfig.CheckAliveness = false; + + clientConfig.RetryPolicy = RetryPolicyForPipes; + t.second.PipeClient = ctx.RegisterWithSameMailbox(NTabletPipe::CreateClient(ctx.SelfID, t.second.TabletID, clientConfig)); + } + + RegisterSessions(ctx); + + ctx.Schedule(Min(CommitInterval, CHECK_ACL_DELAY), new TEvents::TEvWakeup()); + } else { + for (auto& t : ev->Get()->TopicAndTablets) { + if (Topics.find(t.TopicNameConverter->GetInternalName()) == Topics.end()) { + CloseSession(TStringBuilder() << "list of topics changed - new topic '" << + t.TopicNameConverter->GetPrimaryPath() << "' found", + NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + } + } +} + + +void TReadSessionActor::Handle(TEvPersQueue::TEvLockPartition::TPtr& ev, const TActorContext& ctx) { + + auto& record = ev->Get()->Record; + Y_VERIFY(record.GetSession() == Session); + Y_VERIFY(record.GetClientId() == InternalClientId); + + TActorId pipe = ActorIdFromProto(record.GetPipeClient()); + auto path = record.GetPath(); + if (path.empty()) { + path = record.GetTopic(); + } + + auto converterIter = FullPathToConverter.find(NPersQueue::NormalizeFullPath(path)); + if (converterIter.IsEnd()) { + LOG_ALERT_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " ignored ev lock for event = " << record.ShortDebugString() << " path not recognized" + ); + CloseSession( + TStringBuilder() << "Internal server error, cannot parse lock event: " << record.ShortDebugString() << ", reason: topic not found", + NPersQueue::NErrorCode::ERROR, ctx + ); + return; + } + //auto topic = converterIter->second->GetClientsideName(); + auto intName = converterIter->second->GetInternalName(); + Y_VERIFY(!intName.empty()); + auto jt = Topics.find(intName); + + if (jt == Topics.end() || pipe != jt->second.PipeClient) { //this is message from old version of pipe + LOG_DEBUG_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " ignored ev lock for topic = " << converterIter->second->GetPrintableString() + << " path recognized, but topic is unknown, this is unexpected" + ); + return; + } + // ToDo[counters] + if (NumPartitionsFromTopic[intName]++ == 0) { + if (AppData(ctx)->PQConfig.GetTopicsAreFirstClassCitizen()) { + SetupTopicCounters(converterIter->second, jt->second.CloudId, jt->second.DbId, jt->second.FolderId); + } else { + SetupTopicCounters(converterIter->second); + } + } + + auto it = TopicCounters.find(intName); + Y_VERIFY(it != TopicCounters.end()); + + IActor* partitionActor = new TPartitionActor( + ctx.SelfID, InternalClientId, Cookie, Session, record.GetGeneration(), + record.GetStep(), jt->second.FullConverter, record.GetPartition(), record.GetTabletId(), it->second, + ClientDC + ); + + TActorId actorId = ctx.Register(partitionActor); + if (SessionsActive) { + PartsPerSession.DecFor(Partitions.size(), 1); + } + Y_VERIFY(record.GetGeneration() > 0); + //Partitions use clientside name ! + auto pp = Partitions.insert({ + std::make_pair(jt->second.FullConverter->GetClientsideName(), record.GetPartition()), + TPartitionActorInfo{actorId, (((ui64)record.GetGeneration()) << 32) + record.GetStep(), jt->second.FullConverter} + }); + Y_VERIFY(pp.second); + if (SessionsActive) { + PartsPerSession.IncFor(Partitions.size(), 1); + } + + bool res = ActualPartitionActors.insert(actorId).second; + Y_VERIFY(res); + + it->second.PartitionsLocked.Inc(); + it->second.PartitionsInfly.Inc(); + + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " lock: " << record); + + ctx.Send(actorId, new TEvPQProxy::TEvLockPartition(0, 0, false, !ClientsideLocksAllowed)); +} + +void TReadSessionActor::Handle(TEvPQProxy::TEvPartitionStatus::TPtr& ev, const TActorContext&) { + if (!ActualPartitionActor(ev->Sender)) + return; + + auto& evTopic = ev->Get()->Topic; + auto it = Partitions.find(std::make_pair(evTopic->GetClientsideName(), ev->Get()->Partition)); + Y_VERIFY(it != Partitions.end()); + Y_VERIFY(it->second.LockGeneration); + + if (it->second.Releasing) //lock request for already released partition - ignore + return; + + if (ev->Get()->Init) { + Y_VERIFY(!it->second.LockSent); + + it->second.LockSent = true; + auto topicIter = Topics.find(evTopic->GetInternalName()); + Y_VERIFY(topicIter != Topics.end()); + Y_VERIFY(ClientsideLocksAllowed); + TReadResponse result; + auto lock = result.MutableLock(); + lock->SetTopic(topicIter->second.FullConverter->GetClientsideName()); + lock->SetPartition(ev->Get()->Partition); + lock->SetReadOffset(ev->Get()->Offset); + lock->SetEndOffset(ev->Get()->EndOffset); + lock->SetGeneration(it->second.LockGeneration); + auto jt = PartitionToReadResponse.find(it->second.Actor); + if (jt == PartitionToReadResponse.end()) { + ui64 diff = result.ByteSize(); + BytesInflight_ += diff; + if (BytesInflight) (*BytesInflight) += diff; + Handler->Reply(result); + } else { + jt->second->ControlMessages.push_back(result); + } + } else { + Y_VERIFY(it->second.LockSent); + TReadResponse result; + auto status = result.MutablePartitionStatus(); + status->SetTopic(ev->Get()->Topic->GetClientsideName()); + status->SetPartition(ev->Get()->Partition); + status->SetEndOffset(ev->Get()->EndOffset); + status->SetGeneration(it->second.LockGeneration); + status->SetCommittedOffset(ev->Get()->Offset); + status->SetWriteWatermarkMs(ev->Get()->WriteTimestampEstimateMs); + auto jt = PartitionToReadResponse.find(it->second.Actor); + if (jt == PartitionToReadResponse.end()) { + ui64 diff = result.ByteSize(); + BytesInflight_ += diff; + if (BytesInflight) (*BytesInflight) += diff; + Handler->Reply(result); + } else { + jt->second->ControlMessages.push_back(result); + } + + } +} + +void TReadSessionActor::Handle(TEvPersQueue::TEvError::TPtr& ev, const TActorContext& ctx) { + CloseSession(ev->Get()->Record.GetDescription(), ev->Get()->Record.GetCode(), ctx); +} + + +void TReadSessionActor::Handle(TEvPersQueue::TEvReleasePartition::TPtr& ev, const TActorContext& ctx) { + auto& record = ev->Get()->Record; + Y_VERIFY(record.GetSession() == Session); + Y_VERIFY(record.GetClientId() == InternalClientId); + auto topic = record.GetPath(); + if (topic.empty()) { + topic = record.GetTopic(); + } + ui32 group = record.HasGroup() ? record.GetGroup() : 0; + + auto converterIter = FullPathToConverter.find(topic); + if (converterIter.IsEnd()) { + LOG_ALERT_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " Failed to parse balancer response: " << record.ShortDebugString()); + CloseSession( + TStringBuilder() << "Internal server error, cannot parse release event: " << record.ShortDebugString() << ", path not recognized", + NPersQueue::NErrorCode::ERROR, ctx + ); + return; + } + auto name = converterIter->second->GetInternalName(); + + auto it = Topics.find(name); + Y_VERIFY(it != Topics.end()); + + TActorId pipe = ActorIdFromProto(record.GetPipeClient()); + + if (pipe != it->second.PipeClient) { //this is message from old version of pipe + return; + } + + for (ui32 c = 0; c < record.GetCount(); ++c) { + Y_VERIFY(!Partitions.empty()); + + TActorId actorId = TActorId{}; + auto jt = Partitions.begin(); + ui32 i = 0; + for (auto it = Partitions.begin(); it != Partitions.end(); ++it) { + if (it->first.first == name && !it->second.Releasing && (group == 0 || it->first.second + 1 == group)) { + ++i; + if (rand() % i == 0) { //will lead to 1/n probability for each of n partitions + actorId = it->second.Actor; + jt = it; + } + } + } + Y_VERIFY(actorId); + + { + auto it = TopicCounters.find(name); + Y_VERIFY(it != TopicCounters.end()); + it->second.PartitionsToBeReleased.Inc(); + } + + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " releasing " << jt->first.first << ":" << jt->first.second); + jt->second.Releasing = true; + + ctx.Send(actorId, new TEvPQProxy::TEvReleasePartition()); + if (ClientsideLocksAllowed && jt->second.LockSent && !jt->second.Reading) { //locked and no active reads + if (!ProcessReleasePartition(jt, BalanceRightNow, false, ctx)) { // returns false if actor died + return; + } + } + } + AnswerForCommitsIfCan(ctx); // in case of killing partition +} + + +void TReadSessionActor::Handle(TEvPQProxy::TEvPartitionReleased::TPtr& ev, const TActorContext& ctx) { + if (!ActualPartitionActor(ev->Sender)) + return; + + const auto& topic = ev->Get()->Topic; + const ui32 partition = ev->Get()->Partition; + + auto jt = Partitions.find(std::make_pair(topic->GetClientsideName(), partition)); + Y_VERIFY(jt != Partitions.end(), "session %s topic %s part %u", Session.c_str(), topic->GetInternalName().c_str(), partition); + Y_VERIFY(jt->second.Releasing); + jt->second.Released = true; + + { + auto it = TopicCounters.find(topic->GetInternalName()); + Y_VERIFY(it != TopicCounters.end()); + it->second.PartitionsReleased.Inc(); + it->second.PartitionsInfly.Dec(); + it->second.PartitionsToBeReleased.Dec(); + + } + + InformBalancerAboutRelease(jt, ctx); + + DropPartitionIfNeeded(jt, ctx); +} + +void TReadSessionActor::InformBalancerAboutRelease(const THashMap<std::pair<TString, ui32>, TPartitionActorInfo>::iterator& it, const TActorContext& ctx) { + + THolder<TEvPersQueue::TEvPartitionReleased> request; + request.Reset(new TEvPersQueue::TEvPartitionReleased); + auto& req = request->Record; + + auto jt = Topics.find(it->second.Converter->GetInternalName()); + Y_VERIFY(jt != Topics.end()); + + req.SetSession(Session); + ActorIdToProto(jt->second.PipeClient, req.MutablePipeClient()); + req.SetClientId(InternalClientId); + req.SetTopic(it->first.first); + req.SetPartition(it->first.second); + + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " released: " << it->first.first << ":" << it->first.second); + + NTabletPipe::SendData(ctx, jt->second.PipeClient, request.Release()); +} + + +void TReadSessionActor::CloseSession(const TString& errorReason, const NPersQueue::NErrorCode::EErrorCode errorCode, const NActors::TActorContext& ctx) { + + if (errorCode != NPersQueue::NErrorCode::OK) { + + if (InternalErrorCode(errorCode)) { + SLIErrors.Inc(); + } + + if (Errors) { + ++(*Errors); + } else { + if (!AppData(ctx)->PQConfig.GetTopicsAreFirstClassCitizen()) { + ++(*GetServiceCounters(Counters, "pqproxy|readSession")->GetCounter("Errors", true)); + } + } + + TReadResponse result; + + auto error = result.MutableError(); + error->SetDescription(errorReason); + error->SetCode(errorCode); + + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " closed with error reason: " << errorReason); + if (!Handler->IsShuttingDown()) { + ui64 diff = result.ByteSize(); + BytesInflight_ += diff; + if (BytesInflight) (*BytesInflight) += diff; + Handler->Reply(result); + } else { + LOG_WARN_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " GRps is shutting dows, skip reply"); + } + } else { + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " closed"); + } + + Die(ctx); +} + + +void TReadSessionActor::Handle(TEvTabletPipe::TEvClientConnected::TPtr& ev, const TActorContext& ctx) { + TEvTabletPipe::TEvClientConnected *msg = ev->Get(); + if (msg->Status != NKikimrProto::OK) { + if (msg->Dead) { + CloseSession(TStringBuilder() << "one of topics is deleted, tablet " << msg->TabletId, NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + //TODO: remove it + CloseSession(TStringBuilder() << "unable to connect to one of topics, tablet " << msg->TabletId, NPersQueue::NErrorCode::ERROR, ctx); + return; + + const bool isAlive = ProcessBalancerDead(msg->TabletId, ctx); // returns false if actor died + Y_UNUSED(isAlive); + return; + } +} + +bool TReadSessionActor::ActualPartitionActor(const TActorId& part) { + return ActualPartitionActors.contains(part); +} + + +bool TReadSessionActor::ProcessReleasePartition(const THashMap<std::pair<TString, ui32>, TPartitionActorInfo>::iterator& it, + bool kill, bool couldBeReads, const TActorContext& ctx) +{ + //inform client + if (ClientsideLocksAllowed && it->second.LockSent) { + TReadResponse result; + result.MutableRelease()->SetTopic(it->first.first); + result.MutableRelease()->SetPartition(it->first.second); + result.MutableRelease()->SetCanCommit(!kill); + result.MutableRelease()->SetGeneration(it->second.LockGeneration); + auto jt = PartitionToReadResponse.find(it->second.Actor); + if (jt == PartitionToReadResponse.end()) { + ui64 diff = result.ByteSize(); + BytesInflight_ += diff; + if (BytesInflight) (*BytesInflight) += diff; + Handler->Reply(result); + } else { + jt->second->ControlMessages.push_back(result); + } + it->second.LockGeneration = 0; + it->second.LockSent = false; + } + + if (!kill) { + return true; + } + + { + auto jt = TopicCounters.find(it->second.Converter->GetInternalName()); + Y_VERIFY(jt != TopicCounters.end()); + jt->second.PartitionsReleased.Inc(); + jt->second.PartitionsInfly.Dec(); + if (!it->second.Released && it->second.Releasing) { + jt->second.PartitionsToBeReleased.Dec(); + } + } + + //process commits + for (auto& c : it->second.Commits) { + auto kt = Commits.find(c); + Y_VERIFY(kt != Commits.end()); + --kt->second.Partitions; + } + it->second.Commits.clear(); + + Y_VERIFY(couldBeReads || !it->second.Reading); + //process reads + TFormedReadResponse::TPtr formedResponseToAnswer; + if (it->second.Reading) { + const auto readIt = PartitionToReadResponse.find(it->second.Actor); + Y_VERIFY(readIt != PartitionToReadResponse.end()); + if (--readIt->second->RequestsInfly == 0) { + formedResponseToAnswer = readIt->second; + } + } + + InformBalancerAboutRelease(it, ctx); + + it->second.Released = true; //to force drop + DropPartitionIfNeeded(it, ctx); //partition will be dropped + + if (formedResponseToAnswer) { + return ProcessAnswer(ctx, formedResponseToAnswer); // returns false if actor died + } + return true; +} + + +bool TReadSessionActor::ProcessBalancerDead(const ui64 tablet, const TActorContext& ctx) { + for (auto& t : Topics) { + if (t.second.TabletID == tablet) { + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " balancer for topic " << t.first << " is dead, restarting all from this topic"); + + //Drop all partitions from this topic + for (auto it = Partitions.begin(); it != Partitions.end();) { + if (it->second.Converter->GetInternalName() == t.first) { //partition from this topic + // kill actor + auto jt = it; + ++it; + if (!ProcessReleasePartition(jt, true, true, ctx)) { // returns false if actor died + return false; + } + } else { + ++it; + } + } + + AnswerForCommitsIfCan(ctx); + + //reconnect pipe + NTabletPipe::TClientConfig clientConfig; + clientConfig.CheckAliveness = false; + clientConfig.RetryPolicy = RetryPolicyForPipes; + t.second.PipeClient = ctx.RegisterWithSameMailbox(NTabletPipe::CreateClient(ctx.SelfID, t.second.TabletID, clientConfig)); + if (InitDone) { + if (PipeReconnects) { + ++(*PipeReconnects); + ++(*Errors); + } + + RegisterSession(t.second.PipeClient, t.first, ctx); + } + } + } + return true; +} + + +void TReadSessionActor::Handle(TEvTabletPipe::TEvClientDestroyed::TPtr& ev, const TActorContext& ctx) { + const bool isAlive = ProcessBalancerDead(ev->Get()->TabletId, ctx); // returns false if actor died + Y_UNUSED(isAlive); +} + +void TReadSessionActor::ProcessAuth(const NPersQueueCommon::TCredentials& auth) { + TString tmp; + Y_PROTOBUF_SUPPRESS_NODISCARD auth.SerializeToString(&tmp); + if (auth.GetCredentialsCase() != NPersQueueCommon::TCredentials::CREDENTIALS_NOT_SET && tmp != AuthStr) { + Auth = auth; + AuthStr = tmp; + ForceACLCheck = true; + } +} + +void TReadSessionActor::Handle(TEvPQProxy::TEvRead::TPtr& ev, const TActorContext& ctx) { + RequestNotChecked = true; + + THolder<TEvPQProxy::TEvRead> event(ev->Release()); + + Handler->ReadyForNextRead(); + + + ProcessAuth(event->Request.GetCredentials()); + event->Request.ClearCredentials(); + + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " got read request: " << event->Request.GetRead() << " with guid: " << event->Guid); + + Reads.emplace_back(event.Release()); + + const bool isAlive = ProcessReads(ctx); // returns false if actor died + Y_UNUSED(isAlive); +} + + +i64 TReadSessionActor::TFormedReadResponse::ApplyResponse(NPersQueue::TReadResponse&& resp) { + Y_VERIFY(resp.GetBatchedData().PartitionDataSize() == 1); + Response.MutableBatchedData()->AddPartitionData()->Swap(resp.MutableBatchedData()->MutablePartitionData(0)); + i64 prev = Response.ByteSize(); + std::swap<i64>(prev, ByteSize); + return ByteSize - prev; +} + + +void TReadSessionActor::Handle(TEvPQProxy::TEvReadResponse::TPtr& ev, const TActorContext& ctx) { + TActorId sender = ev->Sender; + if (!ActualPartitionActor(sender)) + return; + + THolder<TEvPQProxy::TEvReadResponse> event(ev->Release()); + + Y_VERIFY(event->Response.GetBatchedData().GetCookie() == 0); // cookie is not assigned + Y_VERIFY(event->Response.GetBatchedData().PartitionDataSize() == 1); + + const TString topic = event->Response.GetBatchedData().GetPartitionData(0).GetTopic(); + const ui32 partition = event->Response.GetBatchedData().GetPartitionData(0).GetPartition(); + std::pair<TString, ui32> key(topic, partition); + // Topic is expected to have clientSide name + const auto partitionIt = Partitions.find(key); + Y_VERIFY(partitionIt != Partitions.end()); + Y_VERIFY(partitionIt->second.Reading); + partitionIt->second.Reading = false; + + auto it = PartitionToReadResponse.find(sender); + Y_VERIFY(it != PartitionToReadResponse.end()); + + TFormedReadResponse::TPtr formedResponse = it->second; + + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " read done guid " << formedResponse->Guid + << " " << key.first << ":" << key.second + << " size " << event->Response.ByteSize()); + + const i64 diff = formedResponse->ApplyResponse(std::move(event->Response)); + if (event->FromDisk) { + formedResponse->FromDisk = true; + } + formedResponse->WaitQuotaTime = Max(formedResponse->WaitQuotaTime, event->WaitQuotaTime); + --formedResponse->RequestsInfly; + formedResponse->Offsets.PartitionOffsets.emplace_back(sender, topic, partition, event->NextReadOffset); + + BytesInflight_ += diff; + if (BytesInflight) (*BytesInflight) += diff; + + if (ClientsideLocksAllowed && partitionIt->second.LockSent && partitionIt->second.Releasing) { //locked and need to be released + if (!ProcessReleasePartition(partitionIt, BalanceRightNow, false, ctx)) { // returns false if actor died + return; + } + } + AnswerForCommitsIfCan(ctx); // in case of killing partition + + if (formedResponse->RequestsInfly == 0) { + const bool isAlive = ProcessAnswer(ctx, formedResponse); // returns false if actor died + Y_UNUSED(isAlive); + } +} + + +bool TReadSessionActor::ProcessAnswer(const TActorContext& ctx, TFormedReadResponse::TPtr formedResponse) { + ui32 readDurationMs = (ctx.Now() - formedResponse->Start - formedResponse->WaitQuotaTime).MilliSeconds(); + if (formedResponse->FromDisk) { + ReadLatencyFromDisk.IncFor(readDurationMs, 1); + } else { + ReadLatency.IncFor(readDurationMs, 1); + } + if (readDurationMs >= (formedResponse->FromDisk ? AppData(ctx)->PQConfig.GetReadLatencyFromDiskBigMs() : AppData(ctx)->PQConfig.GetReadLatencyBigMs())) { + SLIBigReadLatency.Inc(); + } + + Y_VERIFY(formedResponse->RequestsInfly == 0); + i64 diff = formedResponse->Response.ByteSize(); + const bool hasMessages = RemoveEmptyMessages(*formedResponse->Response.MutableBatchedData()); + if (hasMessages) { + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " assign read id " << ReadIdToResponse << " to read request " << formedResponse->Guid); + formedResponse->Response.MutableBatchedData()->SetCookie(ReadIdToResponse); + // reply to client + if (ProtocolVersion < NPersQueue::TReadRequest::Batching) { + ConvertToOldBatch(formedResponse->Response); + } + diff -= formedResponse->Response.ByteSize(); // Bytes will be tracked inside handler + Handler->Reply(formedResponse->Response); + } else { + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " empty read result " << formedResponse->Guid << ", start new reading"); + } + + BytesInflight_ -= diff; + if (BytesInflight) (*BytesInflight) -= diff; + + for (auto& r : formedResponse->ControlMessages) { + ui64 diff = r.ByteSize(); + BytesInflight_ += diff; + if (BytesInflight) (*BytesInflight) += diff; + Handler->Reply(r); + } + + for (const TActorId& p : formedResponse->PartitionsTookPartInRead) { + PartitionToReadResponse.erase(p); + } + + // Bring back available partitions. + // If some partition was removed from partitions container, it is not bad because it will be checked during read processing. + AvailablePartitions.insert(formedResponse->PartitionsBecameAvailable.begin(), formedResponse->PartitionsBecameAvailable.end()); + + formedResponse->Offsets.ReadId = ReadIdToResponse; + + RequestedBytes -= formedResponse->RequestedBytes; + + ReadsInfly--; + + if (hasMessages) { + if (!CommitsDisabled) + Offsets.emplace_back(std::move(formedResponse->Offsets)); // even empty responses are needed for correct offsets commit. + ReadIdToResponse++; + } else { + // process new read + NPersQueue::TReadRequest req; + req.MutableRead(); + Reads.emplace_back(new TEvPQProxy::TEvRead(req, formedResponse->Guid)); // Start new reading request with the same guid + } + + return ProcessReads(ctx); // returns false if actor died +} + + +void TReadSessionActor::Handle(TEvPQProxy::TEvCloseSession::TPtr& ev, const TActorContext& ctx) { + CloseSession(ev->Get()->Reason, ev->Get()->ErrorCode, ctx); +} + +void TReadSessionActor::Handle(V1::TEvPQProxy::TEvCloseSession::TPtr& ev, const TActorContext& ctx) { + CloseSession(ev->Get()->Reason, NErrorCode::EErrorCode(ev->Get()->ErrorCode - 500000), ctx); +} + +ui32 TReadSessionActor::NormalizeMaxReadMessagesCount(ui32 sourceValue) { + ui32 count = Min<ui32>(sourceValue, Max<i32>()); + if (count == 0) { + count = Max<i32>(); + } + return count; +} + +ui32 TReadSessionActor::NormalizeMaxReadSize(ui32 sourceValue) { + ui32 size = Min<ui32>(sourceValue, MAX_READ_SIZE); + if (size == 0) { + size = MAX_READ_SIZE; + } + return size; +} + +ui32 TReadSessionActor::NormalizeMaxReadPartitionsCount(ui32 sourceValue) { + ui32 maxPartitions = sourceValue; + if (maxPartitions == 0) { + maxPartitions = Max<ui32>(); + } + return maxPartitions; +} + +bool TReadSessionActor::CheckAndUpdateReadSettings(const NPersQueue::TReadRequest::TRead& readRequest) { + if (ReadSettingsInited) { // already updated. Check that settings are not changed. + const bool hasSettings = readRequest.GetMaxCount() + || readRequest.GetMaxSize() + || readRequest.GetPartitionsAtOnce() + || readRequest.GetMaxTimeLagMs() + || readRequest.GetReadTimestampMs(); + if (!hasSettings) { + return true; + } + + const bool settingsChanged = NormalizeMaxReadMessagesCount(readRequest.GetMaxCount()) != MaxReadMessagesCount + || NormalizeMaxReadSize(readRequest.GetMaxSize()) != MaxReadSize + || NormalizeMaxReadPartitionsCount(readRequest.GetPartitionsAtOnce()) != MaxReadPartitionsCount + || readRequest.GetMaxTimeLagMs() != MaxTimeLagMs + || readRequest.GetReadTimestampMs() != ReadTimestampMs; + return !settingsChanged; + } else { + // Update settings for the first time + ReadSettingsInited = true; + MaxReadMessagesCount = NormalizeMaxReadMessagesCount(readRequest.GetMaxCount()); + MaxReadSize = NormalizeMaxReadSize(readRequest.GetMaxSize()); + MaxReadPartitionsCount = NormalizeMaxReadPartitionsCount(readRequest.GetPartitionsAtOnce()); + MaxTimeLagMs = readRequest.GetMaxTimeLagMs(); + ReadTimestampMs = readRequest.GetReadTimestampMs(); + return true; + } +} + +bool TReadSessionActor::ProcessReads(const TActorContext& ctx) { + while (!Reads.empty() && BytesInflight_ + RequestedBytes < MAX_INFLY_BYTES && ReadsInfly < MAX_INFLY_READS) { + const auto& readRequest = Reads.front()->Request.GetRead(); + if (!CheckAndUpdateReadSettings(readRequest)) { + CloseSession("read settings were changed in read request", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return false; + } + + if (Offsets.size() >= AppData(ctx)->PQConfig.GetMaxReadCookies() + 10) { + CloseSession(TStringBuilder() << "got more than " << AppData(ctx)->PQConfig.GetMaxReadCookies() << " uncommitted reads", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return false; + } + + ui32 count = MaxReadMessagesCount; + ui64 size = MaxReadSize; + ui32 maxPartitions = MaxReadPartitionsCount; + ui32 partitionsAsked = 0; + + TFormedReadResponse::TPtr formedResponse = new TFormedReadResponse(Reads.front()->Guid, ctx.Now()); + while (!AvailablePartitions.empty()) { + auto part = *AvailablePartitions.begin(); + AvailablePartitions.erase(AvailablePartitions.begin()); + + auto it = Partitions.find(std::make_pair(part.Topic->GetClientsideName(), part.Partition)); + if (it == Partitions.end() || it->second.Releasing || it->second.Actor != part.Actor) { //this is already released partition + continue; + } + //add this partition to reading + ++partitionsAsked; + + TAutoPtr<TEvPQProxy::TEvRead> read = new TEvPQProxy::TEvRead(Reads.front()->Request, Reads.front()->Guid); + const ui32 ccount = Min<ui32>(part.MsgLag * LAG_GROW_MULTIPLIER, count); + count -= ccount; + const ui64 csize = (ui64)Min<double>(part.SizeLag * LAG_GROW_MULTIPLIER, size); + size -= csize; + + Y_VERIFY(csize < Max<i32>()); + auto* readR = read->Request.MutableRead(); + readR->SetMaxCount(ccount); + readR->SetMaxSize(csize); + readR->SetMaxTimeLagMs(MaxTimeLagMs); + readR->SetReadTimestampMs(ReadTimestampMs); + + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX + << " performing read request: " << (*readR) << " with guid " << read->Guid + << " from " << part.Topic->GetPrintableString() << ", partition:" << part.Partition + << " count " << ccount << " size " << csize + << " partitionsAsked " << partitionsAsked << " maxTimeLag " << MaxTimeLagMs << "ms"); + + + Y_VERIFY(!it->second.Reading); + it->second.Reading = true; + formedResponse->PartitionsTookPartInRead.insert(it->second.Actor); + + RequestedBytes += csize; + formedResponse->RequestedBytes += csize; + + ctx.Send(it->second.Actor, read.Release()); + const auto insertResult = PartitionToReadResponse.insert(std::make_pair(it->second.Actor, formedResponse)); + Y_VERIFY(insertResult.second); + + if (--maxPartitions == 0 || count == 0 || size == 0) + break; + } + if (partitionsAsked == 0) + break; + ReadsTotal.Inc(); + formedResponse->RequestsInfly = partitionsAsked; + + ReadsInfly++; + + i64 diff = formedResponse->Response.ByteSize(); + BytesInflight_ += diff; + formedResponse->ByteSize = diff; + if (BytesInflight) (*BytesInflight) += diff; + Reads.pop_front(); + } + return true; +} + + +void TReadSessionActor::Handle(TEvPQProxy::TEvPartitionReady::TPtr& ev, const TActorContext& ctx) { + + if (!ActualPartitionActor(ev->Sender)) + return; + + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << ev->Get()->Topic->GetPrintableString() + << " partition:" << ev->Get()->Partition << " ready for read with readOffset " + << ev->Get()->ReadOffset << " endOffset " << ev->Get()->EndOffset << " WTime " + << ev->Get()->WTime << " sizeLag " << ev->Get()->SizeLag); + + const auto it = PartitionToReadResponse.find(ev->Sender); // check whether this partition is taking part in read response + auto& container = it != PartitionToReadResponse.end() ? it->second->PartitionsBecameAvailable : AvailablePartitions; + auto res = container.insert({ev->Get()->Topic, ev->Get()->Partition, ev->Get()->WTime, ev->Get()->SizeLag, + ev->Get()->EndOffset - ev->Get()->ReadOffset, ev->Sender}); + Y_VERIFY(res.second); + const bool isAlive = ProcessReads(ctx); // returns false if actor died + Y_UNUSED(isAlive); +} + + +void TReadSessionActor::HandlePoison(TEvPQProxy::TEvDieCommand::TPtr& ev, const TActorContext& ctx) { + CloseSession(ev->Get()->Reason, ev->Get()->ErrorCode, ctx); +} + + +void TReadSessionActor::HandleWakeup(const TActorContext& ctx) { + ctx.Schedule(Min(CommitInterval, CHECK_ACL_DELAY), new TEvents::TEvWakeup()); + MakeCommit(ctx); + if (!AuthInflight && (ForceACLCheck || (ctx.Now() - LastACLCheckTimestamp > TDuration::Seconds(AppData(ctx)->PQConfig.GetACLRetryTimeoutSec()) && RequestNotChecked))) { + ForceACLCheck = false; + RequestNotChecked = false; + Y_VERIFY(!AuthInitActor); + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " checking auth because of timeout"); + + SendAuthRequest(ctx); + } +} + +bool TReadSessionActor::RemoveEmptyMessages(TReadResponse::TBatchedData& data) { + bool hasNonEmptyMessages = false; + auto isMessageEmpty = [&](TReadResponse::TBatchedData::TMessageData& message) -> bool { + if (message.GetData().empty()) { + return true; + } else { + hasNonEmptyMessages = true; + return false; + } + }; + auto batchRemover = [&](TReadResponse::TBatchedData::TBatch& batch) -> bool { + NProtoBuf::RemoveRepeatedFieldItemIf(batch.MutableMessageData(), isMessageEmpty); + return batch.MessageDataSize() == 0; + }; + auto partitionDataRemover = [&](TReadResponse::TBatchedData::TPartitionData& partition) -> bool { + NProtoBuf::RemoveRepeatedFieldItemIf(partition.MutableBatch(), batchRemover); + return partition.BatchSize() == 0; + }; + NProtoBuf::RemoveRepeatedFieldItemIf(data.MutablePartitionData(), partitionDataRemover); + return hasNonEmptyMessages; +} + + +////////////////// PARTITION ACTOR + +TPartitionActor::TPartitionActor( + const TActorId& parentId, const TString& internalClientId, const ui64 cookie, const TString& session, + const ui32 generation, const ui32 step, const NPersQueue::TTopicConverterPtr& topic, const ui32 partition, + const ui64 tabletID, const TReadSessionActor::TTopicCounters& counters, const TString& clientDC +) + : ParentId(parentId) + , InternalClientId(internalClientId) + , ClientDC(clientDC) + , Cookie(cookie) + , Session(session) + , Generation(generation) + , Step(step) + , Topic(topic) + , Partition(partition) + , TabletID(tabletID) + , ReadOffset(0) + , ClientReadOffset(0) + , ClientCommitOffset(0) + , ClientVerifyReadOffset(false) + , CommittedOffset(0) + , WriteTimestampEstimateMs(0) + , WTime(0) + , InitDone(false) + , StartReading(false) + , AllPrepareInited(false) + , FirstInit(true) + , PipeClient() + , PipeGeneration(0) + , RequestInfly(false) + , EndOffset(0) + , SizeLag(0) + , NeedRelease(false) + , Released(false) + , WaitDataCookie(0) + , WaitForData(false) + , LockCounted(false) + , Counters(counters) +{ +} + + +TPartitionActor::~TPartitionActor() = default; + + +void TPartitionActor::Bootstrap(const TActorContext&) { + Become(&TThis::StateFunc); +} + + +void TPartitionActor::CheckRelease(const TActorContext& ctx) { + const bool hasUncommittedData = ReadOffset > ClientCommitOffset && ReadOffset > ClientReadOffset; + if (NeedRelease) { + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " checking release readOffset " << ReadOffset << " committedOffset " << CommittedOffset << " ReadGuid " << ReadGuid + << " CommitsInfly.size " << CommitsInfly.size() << " Released " << Released); + } + + if (NeedRelease && ReadGuid.empty() && CommitsInfly.empty() && !hasUncommittedData && !Released) { + Released = true; + ctx.Send(ParentId, new TEvPQProxy::TEvPartitionReleased(Topic, Partition)); + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " check release done - releasing; readOffset " << ReadOffset << " committedOffset " << CommittedOffset << " ReadGuid " << ReadGuid + << " CommitsInfly.size " << CommitsInfly.size() << " Released " << Released); + + } +} + + +void TPartitionActor::SendCommit(const ui64 readId, const ui64 offset, const TActorContext& ctx) { + NKikimrClient::TPersQueueRequest request; + request.MutablePartitionRequest()->SetTopic(Topic->GetClientsideName()); + request.MutablePartitionRequest()->SetPartition(Partition); + request.MutablePartitionRequest()->SetCookie(readId); + + Y_VERIFY(PipeClient); + + ActorIdToProto(PipeClient, request.MutablePartitionRequest()->MutablePipeClient()); + auto commit = request.MutablePartitionRequest()->MutableCmdSetClientOffset(); + commit->SetClientId(InternalClientId); + commit->SetOffset(offset); + Y_VERIFY(!Session.empty()); + commit->SetSessionId(Session); + + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" + << Partition << " committing to position " << offset << " prev " << CommittedOffset + << " end " << EndOffset << " by cookie " << readId); + + TAutoPtr<TEvPersQueue::TEvRequest> req(new TEvPersQueue::TEvRequest); + req->Record.Swap(&request); + + NTabletPipe::SendData(ctx, PipeClient, req.Release()); +} + +void TPartitionActor::RestartPipe(const TActorContext& ctx, const TString& reason, const NPersQueue::NErrorCode::EErrorCode errorCode) { + + if (!PipeClient) + return; + + Counters.Errors.Inc(); + + NTabletPipe::CloseClient(ctx, PipeClient); + PipeClient = TActorId{}; + if (errorCode != NPersQueue::NErrorCode::OVERLOAD) + ++PipeGeneration; + + if (PipeGeneration == MAX_PIPE_RESTARTS) { + ctx.Send(ParentId, new TEvPQProxy::TEvCloseSession("too much attempts to restart pipe", NPersQueue::NErrorCode::ERROR)); + return; + } + + ctx.Schedule(TDuration::MilliSeconds(RESTART_PIPE_DELAY_MS), new TEvPQProxy::TEvRestartPipe()); + + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " schedule pipe restart attempt " << PipeGeneration << " reason: " << reason); +} + + +void TPartitionActor::Handle(const TEvPQProxy::TEvRestartPipe::TPtr&, const TActorContext& ctx) { + + Y_VERIFY(!PipeClient); + + NTabletPipe::TClientConfig clientConfig; + clientConfig.RetryPolicy = { + .RetryLimitCount = 6, + .MinRetryTime = TDuration::MilliSeconds(10), + .MaxRetryTime = TDuration::MilliSeconds(100), + .BackoffMultiplier = 2, + .DoFirstRetryInstantly = true + }; + PipeClient = ctx.RegisterWithSameMailbox(NTabletPipe::CreateClient(ctx.SelfID, TabletID, clientConfig)); + Y_VERIFY(TabletID); + + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " pipe restart attempt " << PipeGeneration << " RequestInfly " << RequestInfly << " ReadOffset " << ReadOffset << " EndOffset " << EndOffset + << " InitDone " << InitDone << " WaitForData " << WaitForData); + + if (RequestInfly) { //got read infly + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " resend " << CurrentRequest); + + TAutoPtr<TEvPersQueue::TEvRequest> event(new TEvPersQueue::TEvRequest); + event->Record = CurrentRequest; + + ActorIdToProto(PipeClient, event->Record.MutablePartitionRequest()->MutablePipeClient()); + + NTabletPipe::SendData(ctx, PipeClient, event.Release()); + } + if (InitDone) { + for (auto& c : CommitsInfly) { //resend all commits + if (c.second != Max<ui64>()) + SendCommit(c.first, c.second, ctx); + } + if (WaitForData) { //resend wait-for-data requests + WaitDataInfly.clear(); + WaitDataInPartition(ctx); + } + } +} + +void TPartitionActor::Handle(TEvPersQueue::TEvResponse::TPtr& ev, const TActorContext& ctx) { + + if (ev->Get()->Record.HasErrorCode() && ev->Get()->Record.GetErrorCode() != NPersQueue::NErrorCode::OK) { + const auto errorCode = ev->Get()->Record.GetErrorCode(); + if (errorCode == NPersQueue::NErrorCode::WRONG_COOKIE || errorCode == NPersQueue::NErrorCode::BAD_REQUEST) { + Counters.Errors.Inc(); + ctx.Send(ParentId, new TEvPQProxy::TEvCloseSession("status is not ok: " + ev->Get()->Record.GetErrorReason(), ev->Get()->Record.GetErrorCode())); + } else { + RestartPipe(ctx, TStringBuilder() << "status is not ok. Code: " << EErrorCode_Name(errorCode) << ". Reason: " << ev->Get()->Record.GetErrorReason(), errorCode); + } + return; + } + + if (ev->Get()->Record.GetStatus() != NMsgBusProxy::MSTATUS_OK) { //this is incorrect answer, die + Y_VERIFY(!ev->Get()->Record.HasErrorCode()); + Counters.Errors.Inc(); + ctx.Send(ParentId, new TEvPQProxy::TEvCloseSession("status is not ok: " + ev->Get()->Record.GetErrorReason(), NPersQueue::NErrorCode::ERROR)); + return; + } + if (!ev->Get()->Record.HasPartitionResponse()) { //this is incorrect answer, die + Counters.Errors.Inc(); + ctx.Send(ParentId, new TEvPQProxy::TEvCloseSession("empty partition in response", NPersQueue::NErrorCode::ERROR)); + return; + } + + const auto& result = ev->Get()->Record.GetPartitionResponse(); + + if (!result.HasCookie()) { //this is incorrect answer, die + Counters.Errors.Inc(); + ctx.Send(ParentId, new TEvPQProxy::TEvCloseSession("no cookie in response", NPersQueue::NErrorCode::ERROR)); + return; + } + + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() + << " partition:" << Partition + << " initDone " << InitDone << " event " << PartitionResponseToLog(result)); + + + if (!InitDone) { + if (result.GetCookie() != INIT_COOKIE) { + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() + << " partition:" << Partition + << " unwaited response in init with cookie " << result.GetCookie()); + return; + } + Y_VERIFY(RequestInfly); + CurrentRequest.Clear(); + RequestInfly = false; + + Y_VERIFY(result.HasCmdGetClientOffsetResult()); + const auto& resp = result.GetCmdGetClientOffsetResult(); + Y_VERIFY(resp.HasEndOffset()); + EndOffset = resp.GetEndOffset(); + SizeLag = resp.GetSizeLag(); + + ClientCommitOffset = ReadOffset = CommittedOffset = resp.HasOffset() ? resp.GetOffset() : 0; + Y_VERIFY(EndOffset >= CommittedOffset); + + if (resp.HasWriteTimestampMS()) + WTime = resp.GetWriteTimestampMS(); + WriteTimestampEstimateMs = resp.GetWriteTimestampEstimateMS(); + InitDone = true; + PipeGeneration = 0; //reset tries counter - all ok + LOG_INFO_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " INIT DONE " << Topic->GetPrintableString() + << " partition:" << Partition + << " EndOffset " << EndOffset << " readOffset " << ReadOffset << " committedOffset " << CommittedOffset); + + + if (!StartReading) { + ctx.Send(ParentId, new TEvPQProxy::TEvPartitionStatus(Topic, Partition, CommittedOffset, EndOffset, WriteTimestampEstimateMs, true)); + } else { + InitStartReading(ctx); + } + return; + } + + if (!result.HasCmdReadResult()) { //this is commit response + if (CommitsInfly.empty()) { + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() + << " partition:" << Partition + << " unwaited commit-response with cookie " << result.GetCookie() << "; waiting for nothing"); + return; + } + ui64 readId = CommitsInfly.front().first; + + if (result.GetCookie() != readId) { + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() + << " partition:" << Partition + << " unwaited commit-response with cookie " << result.GetCookie() << "; waiting for " << readId); + return; + } + + Counters.Commits.Inc(); + + CommittedOffset = CommitsInfly.front().second; + CommitsInfly.pop_front(); + if (readId != Max<ui64>()) //this readId is reserved for upcommits on client skipping with ClientCommitOffset + ctx.Send(ParentId, new TEvPQProxy::TEvCommitDone(readId, Topic, Partition)); + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() + << " partition:" << Partition + << " commit done to position " << CommittedOffset << " endOffset " << EndOffset << " with cookie " << readId); + + while (!CommitsInfly.empty() && CommitsInfly.front().second == Max<ui64>()) { //this is cookies that have no effect on this partition + readId = CommitsInfly.front().first; + CommitsInfly.pop_front(); + ctx.Send(ParentId, new TEvPQProxy::TEvCommitDone(readId, Topic, Partition)); + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() + << "partition :" << Partition + << " commit done with no effect with cookie " << readId); + } + + CheckRelease(ctx); + PipeGeneration = 0; //reset tries counter - all ok + return; + } + + //This is read + + Y_VERIFY(result.HasCmdReadResult()); + const auto& res = result.GetCmdReadResult(); + + if (result.GetCookie() != (ui64)ReadOffset) { + LOG_DEBUG_S(ctx, NKikimrServices::PQ_READ_PROXY, PQ_LOG_PREFIX << " " << Topic->GetPrintableString() + << "partition :" << Partition + << " unwaited read-response with cookie " << result.GetCookie() << "; waiting for " << ReadOffset << "; current read guid is " << ReadGuid); + return; + } + + Y_VERIFY(res.HasMaxOffset()); + EndOffset = res.GetMaxOffset(); + SizeLag = res.GetSizeLag(); + + const ui64 realReadOffset = res.HasRealReadOffset() ? res.GetRealReadOffset() : 0; + + TReadResponse response; + + auto* data = response.MutableBatchedData(); + auto* partitionData = data->AddPartitionData(); + partitionData->SetTopic(Topic->GetClientsideName()); + partitionData->SetPartition(Partition); + + bool hasOffset = false; + + TReadResponse::TBatchedData::TBatch* currentBatch = nullptr; + for (ui32 i = 0; i < res.ResultSize(); ++i) { + const auto& r = res.GetResult(i); + + WTime = r.GetWriteTimestampMS(); + WriteTimestampEstimateMs = Max(WriteTimestampEstimateMs, WTime); + Y_VERIFY(r.GetOffset() >= ReadOffset); + ReadOffset = r.GetOffset() + 1; + hasOffset = true; + + auto proto(GetDeserializedData(r.GetData())); + if (proto.GetChunkType() != NKikimrPQClient::TDataChunk::REGULAR) { + continue; //TODO - no such chunks must be on prod + } + + Y_VERIFY(!r.GetSourceId().empty()); + if (!NPQ::NSourceIdEncoding::IsValidEncoded(r.GetSourceId())) { + LOG_ERROR_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << "read bad sourceId from topic " << Topic->GetPrintableString() + << " partition:" << Partition + << " offset " << r.GetOffset() << " seqNo " << r.GetSeqNo() << " sourceId '" << r.GetSourceId() << "' ReadGuid " << ReadGuid); + } + TString sourceId = NPQ::NSourceIdEncoding::Decode(r.GetSourceId()); + + if (!currentBatch || currentBatch->GetWriteTimeMs() != r.GetWriteTimestampMS() || currentBatch->GetSourceId() != sourceId) { + // If write time and source id are the same, the rest fields will be the same too. + currentBatch = partitionData->AddBatch(); + currentBatch->SetWriteTimeMs(r.GetWriteTimestampMS()); + currentBatch->SetSourceId(sourceId); + + if (proto.HasMeta()) { + const auto& header = proto.GetMeta(); + if (header.HasServer()) { + auto* item = currentBatch->MutableExtraFields()->AddItems(); + item->SetKey("server"); + item->SetValue(header.GetServer()); + } + if (header.HasFile()) { + auto* item = currentBatch->MutableExtraFields()->AddItems(); + item->SetKey("file"); + item->SetValue(header.GetFile()); + } + if (header.HasIdent()) { + auto* item = currentBatch->MutableExtraFields()->AddItems(); + item->SetKey("ident"); + item->SetValue(header.GetIdent()); + } + if (header.HasLogType()) { + auto* item = currentBatch->MutableExtraFields()->AddItems(); + item->SetKey("logtype"); + item->SetValue(header.GetLogType()); + } + } + + if (proto.HasExtraFields()) { + const auto& map = proto.GetExtraFields(); + for (const auto& kv : map.GetItems()) { + auto* item = currentBatch->MutableExtraFields()->AddItems(); + item->SetKey(kv.GetKey()); + item->SetValue(kv.GetValue()); + } + } + + if (proto.HasIp() && IsUtf(proto.GetIp())) { + currentBatch->SetIp(proto.GetIp()); + } + } + + auto* message = currentBatch->AddMessageData(); + message->SetSeqNo(r.GetSeqNo()); + message->SetCreateTimeMs(r.GetCreateTimestampMS()); + message->SetOffset(r.GetOffset()); + message->SetUncompressedSize(r.GetUncompressedSize()); + if (proto.HasCodec()) { + const auto codec = proto.GetCodec(); + if (codec < Min<int>() || codec > Max<int>() || !NPersQueueCommon::ECodec_IsValid(codec)) { + LOG_ERROR_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << "data chunk (topic " << Topic->GetInternalName() << ", partition " << Partition + << ", offset " << r.GetOffset() << ", seqNo " << r.GetSeqNo() << ", sourceId " + << r.GetSourceId() << ") codec (id " << codec + << ") is not valid NPersQueueCommon::ECodec, loss of data compression codec information" + ); + } + message->SetCodec((NPersQueueCommon::ECodec)proto.GetCodec()); + } + message->SetData(proto.GetData()); + } + + if (!hasOffset) { //no data could be read from paritition at offset ReadOffset - no data in partition at all??? + ReadOffset = Min(Max(ReadOffset + 1, realReadOffset + 1), EndOffset); + } + + CurrentRequest.Clear(); + RequestInfly = false; + + Y_VERIFY(!WaitForData); + + if (EndOffset > ReadOffset) { + ctx.Send(ParentId, new TEvPQProxy::TEvPartitionReady(Topic, Partition, WTime, SizeLag, ReadOffset, EndOffset)); + } else { + WaitForData = true; + if (PipeClient) //pipe will be recreated soon + WaitDataInPartition(ctx); + } + + LOG_DEBUG_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " after read state " << Topic->GetPrintableString() + << " partition:" << Partition + << " EndOffset " << EndOffset << " ReadOffset " << ReadOffset << " ReadGuid " << ReadGuid); + + ReadGuid = TString(); + + auto readResponse = MakeHolder<TEvPQProxy::TEvReadResponse>( + std::move(response), + ReadOffset, + res.GetBlobsFromDisk() > 0, + TDuration::MilliSeconds(res.GetWaitQuotaTimeMs()) + ); + ctx.Send(ParentId, readResponse.Release()); + CheckRelease(ctx); + + PipeGeneration = 0; //reset tries counter - all ok +} + +void TPartitionActor::Handle(TEvTabletPipe::TEvClientConnected::TPtr& ev, const TActorContext& ctx) { + TEvTabletPipe::TEvClientConnected *msg = ev->Get(); + + LOG_INFO_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " pipe restart attempt " << PipeGeneration << " pipe creation result: " << msg->Status); + + if (msg->Status != NKikimrProto::OK) { + RestartPipe(ctx, TStringBuilder() << "pipe to tablet is dead " << msg->TabletId, NPersQueue::NErrorCode::ERROR); + return; + } +} + +void TPartitionActor::Handle(TEvTabletPipe::TEvClientDestroyed::TPtr& ev, const TActorContext& ctx) { + RestartPipe(ctx, TStringBuilder() << "pipe to tablet is dead " << ev->Get()->TabletId, NPersQueue::NErrorCode::ERROR); +} + + +void TPartitionActor::Handle(TEvPQProxy::TEvReleasePartition::TPtr&, const TActorContext& ctx) { + LOG_INFO_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " (partition)releasing " << Topic->GetPrintableString() << " partition:" << Partition + << " ReadOffset " << ReadOffset << " ClientCommitOffset " << ClientCommitOffset + << " CommittedOffst " << CommittedOffset + ); + NeedRelease = true; + CheckRelease(ctx); +} + + +void TPartitionActor::Handle(TEvPQProxy::TEvGetStatus::TPtr&, const TActorContext& ctx) { + ctx.Send(ParentId, new TEvPQProxy::TEvPartitionStatus(Topic, Partition, CommittedOffset, EndOffset, WriteTimestampEstimateMs, false)); +} + + +void TPartitionActor::Handle(TEvPQProxy::TEvLockPartition::TPtr& ev, const TActorContext& ctx) { + ClientReadOffset = ev->Get()->ReadOffset; + ClientCommitOffset = ev->Get()->CommitOffset; + ClientVerifyReadOffset = ev->Get()->VerifyReadOffset; + + if (StartReading) { + Y_VERIFY(ev->Get()->StartReading); //otherwise it is signal from actor, this could not be done + ctx.Send(ParentId, new TEvPQProxy::TEvCloseSession("double partition locking", NPersQueue::NErrorCode::BAD_REQUEST)); + return; + } + + StartReading = ev->Get()->StartReading; + InitLockPartition(ctx); +} + +void TPartitionActor::InitStartReading(const TActorContext& ctx) { + + Y_VERIFY(AllPrepareInited); + Y_VERIFY(!WaitForData); + LOG_INFO_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " Start reading " << Topic->GetPrintableString() << " partition:" << Partition + << " EndOffset " << EndOffset << " readOffset " << ReadOffset << " committedOffset " + << CommittedOffset << " clientCommittedOffset " << ClientCommitOffset + << " clientReadOffset " << ClientReadOffset + ); + + Counters.PartitionsToBeLocked.Dec(); + LockCounted = false; + + ReadOffset = Max<ui64>(CommittedOffset, ClientReadOffset); + + if (ClientVerifyReadOffset) { + if (ClientReadOffset < CommittedOffset) { + ctx.Send(ParentId, new TEvPQProxy::TEvCloseSession(TStringBuilder() + << "trying to read from position that is less than committed: read " << ClientReadOffset << " committed " << CommittedOffset, + NPersQueue::NErrorCode::BAD_REQUEST)); + return; + } + } + + if (ClientCommitOffset > CommittedOffset) { + if (ClientCommitOffset > ReadOffset) { + ctx.Send(ParentId, new TEvPQProxy::TEvCloseSession(TStringBuilder() + << "trying to read from position that is less than provided to commit: read " << ReadOffset << " commit " << ClientCommitOffset, + NPersQueue::NErrorCode::BAD_REQUEST)); + return; + } + if (ClientCommitOffset > EndOffset) { + ctx.Send(ParentId, new TEvPQProxy::TEvCloseSession(TStringBuilder() + << "trying to commit to future: commit " << ClientCommitOffset << " endOffset " << EndOffset, + NPersQueue::NErrorCode::BAD_REQUEST)); + return; + } + Y_VERIFY(CommitsInfly.empty()); + CommitsInfly.push_back(std::pair<ui64, ui64>(Max<ui64>(), ClientCommitOffset)); + if (PipeClient) //pipe will be recreated soon + SendCommit(CommitsInfly.back().first, CommitsInfly.back().second, ctx); + } else { + ClientCommitOffset = CommittedOffset; + } + + if (EndOffset > ReadOffset) { + ctx.Send(ParentId, new TEvPQProxy::TEvPartitionReady(Topic, Partition, WTime, SizeLag, ReadOffset, EndOffset)); + } else { + WaitForData = true; + if (PipeClient) //pipe will be recreated soon + WaitDataInPartition(ctx); + } +} + +void TPartitionActor::InitLockPartition(const TActorContext& ctx) { + if (PipeClient && AllPrepareInited) { + ctx.Send(ParentId, new TEvPQProxy::TEvCloseSession("double partition locking", NPersQueue::NErrorCode::BAD_REQUEST)); + return; + } + if (!LockCounted) { + Counters.PartitionsToBeLocked.Inc(); + LockCounted = true; + } + if (StartReading) + AllPrepareInited = true; + + if (FirstInit) { + Y_VERIFY(!PipeClient); + FirstInit = false; + NTabletPipe::TClientConfig clientConfig; + clientConfig.RetryPolicy = { + .RetryLimitCount = 6, + .MinRetryTime = TDuration::MilliSeconds(10), + .MaxRetryTime = TDuration::MilliSeconds(100), + .BackoffMultiplier = 2, + .DoFirstRetryInstantly = true + }; + PipeClient = ctx.RegisterWithSameMailbox(NTabletPipe::CreateClient(ctx.SelfID, TabletID, clientConfig)); + + NKikimrClient::TPersQueueRequest request; + + request.MutablePartitionRequest()->SetTopic(Topic->GetClientsideName()); + request.MutablePartitionRequest()->SetPartition(Partition); + request.MutablePartitionRequest()->SetCookie(INIT_COOKIE); + + ActorIdToProto(PipeClient, request.MutablePartitionRequest()->MutablePipeClient()); + + auto cmd = request.MutablePartitionRequest()->MutableCmdCreateSession(); + cmd->SetClientId(InternalClientId); + cmd->SetSessionId(Session); + cmd->SetGeneration(Generation); + cmd->SetStep(Step); + + LOG_INFO_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " INITING " << Topic->GetPrintableString() << " partition:" << Partition); + + TAutoPtr<TEvPersQueue::TEvRequest> req(new TEvPersQueue::TEvRequest); + Y_VERIFY(!RequestInfly); + CurrentRequest = request; + RequestInfly = true; + req->Record.Swap(&request); + + NTabletPipe::SendData(ctx, PipeClient, req.Release()); + } else { + Y_VERIFY(StartReading); //otherwise it is double locking from actor, not client - client makes lock always with StartReading == true + Y_VERIFY(InitDone); + InitStartReading(ctx); + } +} + + +void TPartitionActor::WaitDataInPartition(const TActorContext& ctx) { + + if (WaitDataInfly.size() > 1) //already got 2 requests inflight + return; + Y_VERIFY(InitDone); + + Y_VERIFY(PipeClient); + + if (!WaitForData) + return; + + Y_VERIFY(ReadOffset >= EndOffset); + + TAutoPtr<TEvPersQueue::TEvHasDataInfo> event(new TEvPersQueue::TEvHasDataInfo()); + event->Record.SetPartition(Partition); + event->Record.SetOffset(ReadOffset); + event->Record.SetCookie(++WaitDataCookie); + ui64 deadline = (ctx.Now() + WAIT_DATA - WAIT_DELTA).MilliSeconds(); + event->Record.SetDeadline(deadline); + event->Record.SetClientId(InternalClientId); + + LOG_DEBUG_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " wait data in partition inited, cookie " << WaitDataCookie + ); + + NTabletPipe::SendData(ctx, PipeClient, event.Release()); + + ctx.Schedule(PREWAIT_DATA, new TEvents::TEvWakeup()); + + ctx.Schedule(WAIT_DATA, new TEvPQProxy::TEvDeadlineExceeded(WaitDataCookie)); + + WaitDataInfly.insert(WaitDataCookie); +} + +void TPartitionActor::Handle(TEvPersQueue::TEvHasDataInfoResponse::TPtr& ev, const TActorContext& ctx) { + const auto& record = ev->Get()->Record; + + WriteTimestampEstimateMs = Max(WriteTimestampEstimateMs, record.GetWriteTimestampEstimateMS()); + + auto it = WaitDataInfly.find(ev->Get()->Record.GetCookie()); + if (it == WaitDataInfly.end()) { + LOG_DEBUG_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " unwaited response for WaitData " << ev->Get()->Record); + return; + } + WaitDataInfly.erase(it); + if (!WaitForData) + return; + + Counters.WaitsForData.Inc(); + + Y_VERIFY(record.HasEndOffset()); + Y_VERIFY(EndOffset <= record.GetEndOffset()); //end offset could not be changed if no data arrived, but signal will be sended anyway after timeout + Y_VERIFY(ReadOffset >= EndOffset); //otherwise no WaitData were needed + + LOG_DEBUG_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " wait for data done: " << " readOffset " << ReadOffset << " EndOffset " << EndOffset + << " newEndOffset " << record.GetEndOffset() << " commitOffset " << CommittedOffset + << " clientCommitOffset " << ClientCommitOffset << " cookie " << ev->Get()->Record.GetCookie() + ); + + EndOffset = record.GetEndOffset(); + SizeLag = record.GetSizeLag(); + + if (ReadOffset < EndOffset) { + WaitForData = false; + WaitDataInfly.clear(); + ctx.Send(ParentId, new TEvPQProxy::TEvPartitionReady(Topic, Partition, WTime, SizeLag, ReadOffset, EndOffset)); + LOG_DEBUG_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " ready for read with readOffset " << ReadOffset << " endOffset " << EndOffset + ); + } else { + if (PipeClient) + WaitDataInPartition(ctx); + } + CheckRelease(ctx); //just for logging purpose +} + + +void TPartitionActor::Handle(TEvPQProxy::TEvRead::TPtr& ev, const TActorContext& ctx) { + LOG_DEBUG_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " READ FROM " << Topic->GetPrintableString() << " partition:" << Partition + << " event " << ev->Get()->Request << " readOffset " << ReadOffset + << " EndOffset " << EndOffset << " ClientCommitOffset " << ClientCommitOffset + << " committedOffset " << CommittedOffset << " Guid " << ev->Get()->Guid + ); + + Y_VERIFY(!NeedRelease); + Y_VERIFY(!Released); + + Y_VERIFY(ReadGuid.empty()); + Y_VERIFY(!RequestInfly); + + ReadGuid = ev->Get()->Guid; + + const auto& req = ev->Get()->Request.GetRead(); + + NKikimrClient::TPersQueueRequest request; + + request.MutablePartitionRequest()->SetTopic(Topic->GetClientsideName()); + + request.MutablePartitionRequest()->SetPartition(Partition); + request.MutablePartitionRequest()->SetCookie((ui64)ReadOffset); + + ActorIdToProto(PipeClient, request.MutablePartitionRequest()->MutablePipeClient()); + auto read = request.MutablePartitionRequest()->MutableCmdRead(); + read->SetClientId(InternalClientId); + read->SetClientDC(ClientDC); + if (req.GetMaxCount()) { + read->SetCount(req.GetMaxCount()); + } + if (req.GetMaxSize()) { + read->SetBytes(req.GetMaxSize()); + } + if (req.GetMaxTimeLagMs()) { + read->SetMaxTimeLagMs(req.GetMaxTimeLagMs()); + } + if (req.GetReadTimestampMs()) { + read->SetReadTimestampMs(req.GetReadTimestampMs()); + } + + read->SetOffset(ReadOffset); + read->SetTimeoutMs(READ_TIMEOUT_DURATION.MilliSeconds()); + RequestInfly = true; + CurrentRequest = request; + + if (!PipeClient) //Pipe will be recreated soon + return; + + TAutoPtr<TEvPersQueue::TEvRequest> event(new TEvPersQueue::TEvRequest); + event->Record.Swap(&request); + + NTabletPipe::SendData(ctx, PipeClient, event.Release()); +} + + +void TPartitionActor::Handle(TEvPQProxy::TEvCommit::TPtr& ev, const TActorContext& ctx) { + const ui64 readId = ev->Get()->ReadId; + const ui64 offset = ev->Get()->Offset; + Y_VERIFY(offset != Max<ui64>()); // has concreete offset + if (offset < ClientCommitOffset) { + LOG_ERROR_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " commit done to too small position " << offset + << " committedOffset " << ClientCommitOffset << " cookie " << readId + ); + } + Y_VERIFY(offset >= ClientCommitOffset); + + const bool hasProgress = offset > ClientCommitOffset; + + if (!hasProgress) {//nothing to commit for this partition + if (CommitsInfly.empty()) { + LOG_DEBUG_S( + ctx, NKikimrServices::PQ_READ_PROXY, + PQ_LOG_PREFIX << " " << Topic->GetPrintableString() << " partition:" << Partition + << " commit done with no effect with cookie " << readId + ); + ctx.Send(ParentId, new TEvPQProxy::TEvCommitDone(readId, Topic, Partition)); + CheckRelease(ctx); + } else { + CommitsInfly.push_back(std::pair<ui64, ui64>(readId, Max<ui64>())); + } + return; + } + + ClientCommitOffset = offset; + CommitsInfly.push_back(std::pair<ui64, ui64>(readId, offset)); + + if (PipeClient) //if not then pipe will be recreated soon and SendCommit will be done + SendCommit(readId, offset, ctx); +} + + +void TPartitionActor::Die(const TActorContext& ctx) { + if (PipeClient) + NTabletPipe::CloseClient(ctx, PipeClient); + TActorBootstrapped<TPartitionActor>::Die(ctx); +} + +void TPartitionActor::HandlePoison(TEvents::TEvPoisonPill::TPtr&, const TActorContext& ctx) { + if (LockCounted) + Counters.PartitionsToBeLocked.Dec(); + Die(ctx); +} + +void TPartitionActor::Handle(TEvPQProxy::TEvDeadlineExceeded::TPtr& ev, const TActorContext& ctx) { + + WaitDataInfly.erase(ev->Get()->Cookie); + if (ReadOffset >= EndOffset && WaitDataInfly.size() <= 1 && PipeClient) { + Y_VERIFY(WaitForData); + WaitDataInPartition(ctx); + } + +} + +void TPartitionActor::HandleWakeup(const TActorContext& ctx) { + if (ReadOffset >= EndOffset && WaitDataInfly.size() <= 1 && PipeClient) { + Y_VERIFY(WaitForData); + WaitDataInPartition(ctx); + } +} +} +} diff --git a/kikimr/yndx/grpc_services/persqueue/grpc_pq_session.h b/kikimr/yndx/grpc_services/persqueue/grpc_pq_session.h new file mode 100644 index 0000000000..22e2b61e5b --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/grpc_pq_session.h @@ -0,0 +1,317 @@ +#pragma once + +#include "ydb/core/client/server/grpc_base.h" +#include <library/cpp/grpc/server/grpc_server.h> +#include <util/generic/queue.h> + +using grpc::Status; + + +namespace NKikimr { +namespace NGRpcProxy { + +/////////////////////////////////////////////////////////////////////////////// + +using namespace NKikimrClient; + +template<class TResponse> +class ISessionHandler : public TAtomicRefCount<ISessionHandler<TResponse>> { +public: + virtual ~ISessionHandler() + { } + + /// Finish session. + virtual void Finish() = 0; + + /// Send reply to client. + virtual void Reply(const TResponse& resp) = 0; + + virtual void ReadyForNextRead() = 0; + + virtual bool IsShuttingDown() const = 0; +}; + +template<class TResponse> +using ISessionHandlerRef = TIntrusivePtr<ISessionHandler<TResponse>>; + + +template <class TRequest, class TResponse> +class ISession : public ISessionHandler<TResponse> +{ + + using ISessionRef = TIntrusivePtr<ISession<TRequest, TResponse>>; + +protected: + class TRequestCreated : public NGrpc::IQueueEvent { + public: + TRequestCreated(ISessionRef session) + : Session(session) + { } + + bool Execute(bool ok) override { + if (!ok) { + Session->DestroyStream("waiting stream creating failed"); + return false; + } + + Session->OnCreated(); + return false; + } + + void DestroyRequest() override { + if (!Session->Context.c_call() && Session->ClientDone) { + // AsyncNotifyWhenDone will not appear on the queue. + delete Session->ClientDone; + Session->ClientDone = nullptr; + } + delete this; + } + + ISessionRef Session; + }; + + class TReadDone : public NGrpc::IQueueEvent { + public: + TReadDone(ISessionRef session) + : Session(session) + { } + + bool Execute(bool ok) override { + if (ok) { + Session->OnRead(Request); + } else { + if (Session->IsCancelled()) { + Session->DestroyStream("reading from stream failed"); + } else { + Session->OnDone(); + } + } + return false; + } + + void DestroyRequest() override { + delete this; + } + + TRequest Request; + ISessionRef Session; + }; + + class TWriteDone : public NGrpc::IQueueEvent { + public: + TWriteDone(ISessionRef session, ui64 size) + : Session(session) + , Size(size) + { } + + bool Execute(bool ok) override { + Session->OnWriteDone(Size); + if (!ok) { + Session->DestroyStream("writing to stream failed"); + return false; + } + + TGuard<TSpinLock> lock(Session->Lock); + if (Session->Responses.empty()) { + Session->HaveWriteInflight = false; + if (Session->NeedFinish) { + lock.Release(); + Session->Stream.Finish(Status::OK, new TFinishDone(Session)); + } + } else { + auto resp = Session->Responses.front(); + Session->Responses.pop(); + lock.Release(); + ui64 sz = resp.ByteSize(); + Session->Stream.Write(resp, new TWriteDone(Session, sz)); + } + + return false; + } + + void DestroyRequest() override { + delete this; + } + + ISessionRef Session; + ui64 Size; + }; + + class TFinishDone : public NGrpc::IQueueEvent { + public: + TFinishDone(ISessionRef session) + : Session(session) + { } + + bool Execute(bool) override { + Session->DestroyStream("some stream finished"); + return false; + } + + void DestroyRequest() override { + delete this; + } + + ISessionRef Session; + }; + + class TClientDone : public NGrpc::IQueueEvent { + public: + TClientDone(ISessionRef session) + : Session(session) + { + Session->ClientDone = this; + } + + bool Execute(bool) override { + Session->ClientIsDone = true; + Session->DestroyStream("sesison closed"); + return false; + } + + void DestroyRequest() override { + Y_VERIFY(Session->ClientDone); + Session->ClientDone = nullptr; + delete this; + } + + ISessionRef Session; + }; + +public: + ISession(grpc::ServerCompletionQueue* cq) + : CQ(cq) + , Stream(&Context) + , HaveWriteInflight(false) + , NeedFinish(false) + , ClientIsDone(false) + { + Context.AsyncNotifyWhenDone(new TClientDone(this)); + } + + TString GetDatabase() const { + TString key = "x-ydb-database"; + const auto& clientMetadata = Context.client_metadata(); + const auto range = clientMetadata.equal_range(grpc::string_ref{key.data(), key.size()}); + if (range.first == range.second) { + return ""; + } + + TVector<TStringBuf> values; + values.reserve(std::distance(range.first, range.second)); + + for (auto it = range.first; it != range.second; ++it) { + return TString(it->second.data(), it->second.size()); + } + return ""; + } + + TString GetPeerName() const { + TString res(Context.peer()); + if (res.StartsWith("ipv4:[") || res.StartsWith("ipv6:[")) { + size_t pos = res.find(']'); + Y_VERIFY(pos != TString::npos); + res = res.substr(6, pos - 6); + } else if (res.StartsWith("ipv4:")) { + size_t pos = res.rfind(':'); + if (pos == TString::npos) {//no port + res = res.substr(5); + } else { + res = res.substr(5, pos - 5); + } + } else { + size_t pos = res.rfind(":"); //port + if (pos != TString::npos) { + res = res.substr(0, pos); + } + } + return res; + } + +protected: + + virtual void OnCreated() = 0; + virtual void OnRead(const TRequest& request) = 0; + virtual void OnDone() = 0; + virtual void OnWriteDone(ui64 size) = 0; + + virtual void DestroyStream(const TString& reason, NPersQueue::NErrorCode::EErrorCode code = NPersQueue::NErrorCode::BAD_REQUEST) = 0; + + /// Start accepting session's requests. + virtual void Start() = 0; + + bool IsCancelled() const { + return ClientIsDone && Context.IsCancelled(); + } + + void ReplyWithError(const TString& description, NPersQueue::NErrorCode::EErrorCode code) + { + TResponse response; + response.MutableError()->SetDescription(description); + response.MutableError()->SetCode(code); + Reply(response); + Finish(); + } + + /// Finish session. + void Finish() override { + { + TGuard<TSpinLock> lock(Lock); + if (NeedFinish) + return; + if (HaveWriteInflight || !Responses.empty()) { + NeedFinish = true; + return; + } + HaveWriteInflight = true; + } + + Stream.Finish(Status::OK, new TFinishDone(this)); + } + + /// Send reply to client. + void Reply(const TResponse& resp) override { + { + TGuard<TSpinLock> lock(Lock); + if (NeedFinish) //ignore responses after finish + return; + if (HaveWriteInflight || !Responses.empty()) { + Responses.push(resp); + return; + } else { + HaveWriteInflight = true; + } + } + + ui64 size = resp.ByteSize(); + Stream.Write(resp, new TWriteDone(this, size)); + } + + void ReadyForNextRead() override { + { + TGuard<TSpinLock> lock(Lock); + if (NeedFinish) { + return; + } + } + + auto read = new TReadDone(this); + Stream.Read(&read->Request, read); + } + +protected: + grpc::ServerCompletionQueue* const CQ; + grpc::ServerContext Context; + grpc::ServerAsyncReaderWriter<TResponse, TRequest> + Stream; +private: + TSpinLock Lock; + bool HaveWriteInflight; + bool NeedFinish; + std::atomic<bool> ClientIsDone; + TClientDone* ClientDone; + TQueue<TResponse> Responses; //TODO: if Responses total size is too big - fail this session; +}; + +} +} diff --git a/kikimr/yndx/grpc_services/persqueue/grpc_pq_write.cpp b/kikimr/yndx/grpc_services/persqueue/grpc_pq_write.cpp new file mode 100644 index 0000000000..36ba3fa8f6 --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/grpc_pq_write.cpp @@ -0,0 +1,221 @@ +#include "grpc_pq_write.h" +#include "grpc_pq_actor.h" +#include "grpc_pq_session.h" +#include "ydb/core/client/server/grpc_proxy_status.h" + +#include <ydb/core/base/appdata.h> +#include <util/generic/queue.h> + +using namespace NActors; +using namespace NKikimrClient; + +using grpc::Status; + +namespace NKikimr { +namespace NGRpcProxy { + +using namespace NPersQueue; + +/////////////////////////////////////////////////////////////////////////////// + + +void TPQWriteServiceImpl::TSession::OnCreated() { // Start waiting for new session. + Proxy->WaitWriteSession(); + if (Proxy->TooMuchSessions()) { + ReplyWithError("proxy overloaded", NPersQueue::NErrorCode::OVERLOAD); + return; + } + TMaybe<TString> localCluster = Proxy->AvailableLocalCluster(); + if (NeedDiscoverClusters) { + if (!localCluster.Defined()) { + ReplyWithError("initializing", NPersQueue::NErrorCode::INITIALIZING); + return; + } else if (localCluster->empty()) { + ReplyWithError("cluster disabled", NPersQueue::NErrorCode::CLUSTER_DISABLED); + return; + } else { + CreateActor(*localCluster); + } + } else { + CreateActor(TString()); + } + ReadyForNextRead(); +} + +void TPQWriteServiceImpl::TSession::OnRead(const TWriteRequest& request) { + + switch (request.GetRequestCase()) { + case TWriteRequest::kInit: { + SendEvent(new TEvPQProxy::TEvWriteInit(request, GetPeerName(), GetDatabase())); + break; + } + case TWriteRequest::kDataBatch: + case TWriteRequest::kData: { + SendEvent(new TEvPQProxy::TEvWrite(request)); + break; + } + default: { + ReplyWithError("unsupported request", NPersQueue::NErrorCode::BAD_REQUEST); + } + } +} + +void TPQWriteServiceImpl::TSession::OnDone() { + SendEvent(new TEvPQProxy::TEvDone()); +} + +TPQWriteServiceImpl::TSession::TSession(std::shared_ptr<TPQWriteServiceImpl> proxy, + grpc::ServerCompletionQueue* cq, ui64 cookie, const TActorId& schemeCache, + TIntrusivePtr<NMonitoring::TDynamicCounters> counters, bool needDiscoverClusters) + : ISession(cq) + , Proxy(proxy) + , Cookie(cookie) + , SchemeCache(schemeCache) + , Counters(counters) + , NeedDiscoverClusters(needDiscoverClusters) +{ +} + +void TPQWriteServiceImpl::TSession::Start() { + if (!Proxy->IsShuttingDown()) { + Proxy->RequestSession(&Context, &Stream, CQ, CQ, new TRequestCreated(this)); + } +} + +ui64 TPQWriteServiceImpl::TSession::GetCookie() const { + return Cookie; +} + +void TPQWriteServiceImpl::TSession::DestroyStream(const TString& reason, const NPersQueue::NErrorCode::EErrorCode errorCode) { + // Send poison pill to the actor(if it is alive) + SendEvent(new TEvPQProxy::TEvDieCommand("write-session " + ToString<ui64>(Cookie) + ": " + reason, errorCode)); + // Remove reference to session from "cookie -> session" map. + Proxy->ReleaseSession(this); +} + +bool TPQWriteServiceImpl::TSession::IsShuttingDown() const { + return Proxy->IsShuttingDown(); +} + +void TPQWriteServiceImpl::TSession::CreateActor(const TString &localCluster) { + + auto classifier = Proxy->GetClassifier(); + ActorId = Proxy->ActorSystem->Register( + new TWriteSessionActor(this, Cookie, SchemeCache, Counters, localCluster, + classifier ? classifier->ClassifyAddress(GetPeerName()) + : "unknown"), TMailboxType::Simple, 0 + ); +} + +void TPQWriteServiceImpl::TSession::SendEvent(IEventBase* ev) { + Proxy->ActorSystem->Send(ActorId, ev); +} + +/////////////////////////////////////////////////////////////////////////////// + + +TPQWriteServiceImpl::TPQWriteServiceImpl(grpc::ServerCompletionQueue* cq, + NActors::TActorSystem* as, const TActorId& schemeCache, + TIntrusivePtr<NMonitoring::TDynamicCounters> counters, const ui32 maxSessions) + : CQ(cq) + , ActorSystem(as) + , SchemeCache(schemeCache) + , Counters(counters) + , MaxSessions(maxSessions) + , NeedDiscoverClusters(false) +{ +} + +void TPQWriteServiceImpl::InitClustersUpdater() +{ + TAppData* appData = ActorSystem->AppData<TAppData>(); + NeedDiscoverClusters = !appData->PQConfig.GetTopicsAreFirstClassCitizen(); + if (NeedDiscoverClusters) { + ActorSystem->Register(new TClustersUpdater(this)); + } +} + + +ui64 TPQWriteServiceImpl::NextCookie() { + return AtomicIncrement(LastCookie); +} + + +void TPQWriteServiceImpl::ReleaseSession(TSessionRef session) { + with_lock (Lock) { + bool erased = Sessions.erase(session->GetCookie()); + if (erased) { + ActorSystem->Send(MakeGRpcProxyStatusID(ActorSystem->NodeId), new TEvGRpcProxyStatus::TEvUpdateStatus(0, 0, -1, 0)); + } + } +} + + +void TPQWriteServiceImpl::SetupIncomingRequests() { + WaitWriteSession(); +} + + +void TPQWriteServiceImpl::WaitWriteSession() { + + const ui64 cookie = NextCookie(); + + ActorSystem->Send(MakeGRpcProxyStatusID(ActorSystem->NodeId), new TEvGRpcProxyStatus::TEvUpdateStatus(0,0,1,0)); + + TSessionRef session(new TSession(shared_from_this(), CQ, cookie, SchemeCache, Counters, NeedDiscoverClusters)); + { + with_lock (Lock) { + Sessions[cookie] = session; + } + } + + session->Start(); +} + + +bool TPQWriteServiceImpl::TooMuchSessions() { + with_lock (Lock) { + return Sessions.size() >= MaxSessions; + } +} + + +TMaybe<TString> TPQWriteServiceImpl::AvailableLocalCluster() { + with_lock (Lock) { + return AvailableLocalClusterName; + } +} + + +void TPQWriteServiceImpl::NetClassifierUpdated(NAddressClassifier::TLabeledAddressClassifier::TConstPtr classifier) { + auto g(Guard(Lock)); + if (!DatacenterClassifier) { + for (auto it = Sessions.begin(); it != Sessions.end();) { + auto jt = it++; + jt->second->DestroyStream("datacenter classifier initialized, restart session please", NPersQueue::NErrorCode::INITIALIZING); + } + } + + DatacenterClassifier = classifier; +} + + + +void TPQWriteServiceImpl::CheckClusterChange(const TString &localCluster, const bool enabled) { + with_lock (Lock) { + AvailableLocalClusterName = enabled ? localCluster : TString(); + + if (!enabled) { + for (auto it = Sessions.begin(); it != Sessions.end();) { + auto jt = it++; + jt->second->DestroyStream("cluster disabled", NPersQueue::NErrorCode::CLUSTER_DISABLED); + } + } + } +} + + +/////////////////////////////////////////////////////////////////////////////// + +} +} diff --git a/kikimr/yndx/grpc_services/persqueue/grpc_pq_write.h b/kikimr/yndx/grpc_services/persqueue/grpc_pq_write.h new file mode 100644 index 0000000000..40a0d64e47 --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/grpc_pq_write.h @@ -0,0 +1,148 @@ +#pragma once + +#include "grpc_pq_clusters_updater_actor.h" +#include "grpc_pq_session.h" + +#include <ydb/core/client/server/grpc_base.h> + +#include <kikimr/yndx/api/grpc/persqueue.grpc.pb.h> + +#include <library/cpp/actors/core/actorsystem.h> + +#include <util/generic/hash.h> +#include <util/generic/maybe.h> +#include <util/system/mutex.h> + +namespace NKikimr { +namespace NGRpcProxy { + +// Класс, отвечающий за обработку запросов на запись. + +class TPQWriteServiceImpl : public IPQClustersUpdaterCallback, public std::enable_shared_from_this<TPQWriteServiceImpl> { + + class TSession : public ISession<NPersQueue::TWriteRequest, NPersQueue::TWriteResponse> + { + + void OnCreated() override; + void OnRead(const NPersQueue::TWriteRequest& request) override; + void OnDone() override; + void OnWriteDone(ui64) override {}; + + public: + TSession(std::shared_ptr<TPQWriteServiceImpl> proxy, + grpc::ServerCompletionQueue* cq, ui64 cookie, const TActorId& schemeCache, + TIntrusivePtr<NMonitoring::TDynamicCounters> counters, bool needDiscoverClusters); + void Start() override; + ui64 GetCookie() const; + void DestroyStream(const TString& reason, const NPersQueue::NErrorCode::EErrorCode errorCode) override; + bool IsShuttingDown() const override; + + private: + void CreateActor(const TString& localCluster); + void SendEvent(NActors::IEventBase* ev); + + private: + std::shared_ptr<TPQWriteServiceImpl> Proxy; + const ui64 Cookie; + + NActors::TActorId ActorId; + + const NActors::TActorId SchemeCache; + + TIntrusivePtr<NMonitoring::TDynamicCounters> Counters; + + bool NeedDiscoverClusters; + }; + using TSessionRef = TIntrusivePtr<TSession>; + +public: + TPQWriteServiceImpl(grpc::ServerCompletionQueue* cq, + NActors::TActorSystem* as, const NActors::TActorId& schemeCache, TIntrusivePtr<NMonitoring::TDynamicCounters> counters, + const ui32 maxSessions); + virtual ~TPQWriteServiceImpl() = default; + + void SetupIncomingRequests(); + + virtual void RequestSession(::grpc::ServerContext* context, ::grpc::ServerAsyncReaderWriter< ::NPersQueue::TWriteResponse, ::NPersQueue::TWriteRequest>* stream, + ::grpc::CompletionQueue* new_call_cq, ::grpc::ServerCompletionQueue* notification_cq, void *tag) = 0; + + void StopService() { + AtomicSet(ShuttingDown_, 1); + } + + bool IsShuttingDown() const { + return AtomicGet(ShuttingDown_); + } + void InitClustersUpdater(); + +private: + ui64 NextCookie(); + + //! Unregistry session object. + void ReleaseSession(TSessionRef session); + + //! Start listening for incoming connections. + void WaitWriteSession(); + bool TooMuchSessions(); + TMaybe<TString> AvailableLocalCluster(); + NAddressClassifier::TLabeledAddressClassifier::TConstPtr GetClassifier() const { + auto g(Guard(Lock)); + return DatacenterClassifier; + } + void CheckClusterChange(const TString& localCluster, const bool enabled) override; + void NetClassifierUpdated(NAddressClassifier::TLabeledAddressClassifier::TConstPtr classifier) override; + +private: + grpc::ServerContext Context; + grpc::ServerCompletionQueue* CQ; + + NActors::TActorSystem* ActorSystem; + NActors::TActorId SchemeCache; + + TAtomic LastCookie = 0; + + TMutex Lock; + THashMap<ui64, TSessionRef> Sessions; + + TIntrusivePtr<NMonitoring::TDynamicCounters> Counters; + + ui32 MaxSessions; + TMaybe<TString> AvailableLocalClusterName; + TString SelectSourceIdQuery; + TString UpdateSourceIdQuery; + TString DeleteSourceIdQuery; + + TAtomic ShuttingDown_ = 0; + + bool NeedDiscoverClusters; // Legacy mode OR account-mode in multi-cluster setup; + + NAddressClassifier::TLabeledAddressClassifier::TConstPtr DatacenterClassifier; // Detects client's datacenter by IP. May be null +}; + + +class TPQWriteService : public TPQWriteServiceImpl { +public: + TPQWriteService(NPersQueue::PersQueueService::AsyncService* service, + grpc::ServerCompletionQueue* cq, + NActors::TActorSystem* as, const NActors::TActorId& schemeCache, TIntrusivePtr<NMonitoring::TDynamicCounters> counters, + const ui32 maxSessions) + : TPQWriteServiceImpl(cq, as, schemeCache, counters, maxSessions) + , Service(service) + {} + + virtual ~TPQWriteService() + {} + + void RequestSession(::grpc::ServerContext* context, ::grpc::ServerAsyncReaderWriter< ::NPersQueue::TWriteResponse, ::NPersQueue::TWriteRequest>* stream, + ::grpc::CompletionQueue* new_call_cq, ::grpc::ServerCompletionQueue* notification_cq, void *tag) override + { + Service->RequestWriteSession(context, stream, new_call_cq, notification_cq, tag); + } + +private: + NPersQueue::PersQueueService::AsyncService* Service; +}; + + +} +} diff --git a/kikimr/yndx/grpc_services/persqueue/grpc_pq_write_actor.cpp b/kikimr/yndx/grpc_services/persqueue/grpc_pq_write_actor.cpp new file mode 100644 index 0000000000..ae2c3198c0 --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/grpc_pq_write_actor.cpp @@ -0,0 +1,1055 @@ +#include "grpc_pq_actor.h" + +#include <ydb/core/persqueue/pq_database.h> +#include <ydb/core/persqueue/write_meta.h> +#include <ydb/core/protos/services.pb.h> +#include <ydb/public/lib/deprecated/kicli/kicli.h> +#include <ydb/library/persqueue/topic_parser/topic_parser.h> +#include <ydb/library/persqueue/topic_parser/counters.h> +#include <ydb/services/lib/sharding/sharding.h> + +#include <library/cpp/actors/core/log.h> +#include <library/cpp/digest/md5/md5.h> +#include <util/string/hex.h> +#include <util/string/vector.h> +#include <util/string/escape.h> + +using namespace NActors; +using namespace NKikimrClient; + + +namespace NKikimr { +using namespace NMsgBusProxy::NPqMetaCacheV2; +using namespace NSchemeCache; +using namespace NPQ; + +template <> +void FillChunkDataFromReq(NKikimrPQClient::TDataChunk& proto, const NPersQueue::TWriteRequest::TData& data) { + proto.SetData(data.GetData()); + proto.SetSeqNo(data.GetSeqNo()); + proto.SetCreateTime(data.GetCreateTimeMs()); + proto.SetCodec(data.GetCodec()); +} + +template <> +void FillExtraFieldsForDataChunk( + const NPersQueue::TWriteRequest::TInit& init, + NKikimrPQClient::TDataChunk& data, + TString& server, + TString& ident, + TString& logType, + TString& file +) { + for (ui32 i = 0; i < init.GetExtraFields().ItemsSize(); ++i) { + const auto& item = init.GetExtraFields().GetItems(i); + if (item.GetKey() == "server") { + server = item.GetValue(); + } else if (item.GetKey() == "ident") { + ident = item.GetValue(); + } else if (item.GetKey() == "logtype") { + logType = item.GetValue(); + } else if (item.GetKey() == "file") { + file = item.GetValue(); + } else { + auto res = data.MutableExtraFields()->AddItems(); + res->SetKey(item.GetKey()); + res->SetValue(item.GetValue()); + } + } +} + +namespace NGRpcProxy { + +using namespace NPersQueue; + +static const ui32 MAX_RESERVE_REQUESTS_INFLIGHT = 5; + +static const ui32 MAX_BYTES_INFLIGHT = 1_MB; +static const TDuration SOURCEID_UPDATE_PERIOD = TDuration::Hours(1); + +//TODO: add here tracking of bytes in/out + +TWriteSessionActor::TWriteSessionActor(IWriteSessionHandlerRef handler, const ui64 cookie, const TActorId& schemeCache, + TIntrusivePtr<NMonitoring::TDynamicCounters> counters, const TString& localDC, + const TMaybe<TString> clientDC) + : Handler(handler) + , State(ES_CREATED) + , SchemeCache(schemeCache) + , PeerName("") + , Cookie(cookie) + , Partition(0) + , PreferedPartition(Max<ui32>()) + , NumReserveBytesRequests(0) + , WritesDone(false) + , Counters(counters) + , BytesInflight_(0) + , BytesInflightTotal_(0) + , NextRequestInited(false) + , NextRequestCookie(0) + , Token(nullptr) + , ACLCheckInProgress(true) + , FirstACLCheck(true) + , ForceACLCheck(false) + , RequestNotChecked(true) + , LastACLCheckTimestamp(TInstant::Zero()) + , LogSessionDeadline(TInstant::Zero()) + , BalancerTabletId(0) + , PipeToBalancer() + , LocalDC(localDC) + , ClientDC(clientDC ? *clientDC : "other") + , LastSourceIdUpdate(TInstant::Zero()) + , SourceIdCreateTime(0) + , SourceIdUpdatesInflight(0) + +{ + Y_ASSERT(Handler); +} + +TWriteSessionActor::~TWriteSessionActor() = default; + + +void TWriteSessionActor::Bootstrap(const TActorContext& ctx) { + if (!AppData(ctx)->PQConfig.GetTopicsAreFirstClassCitizen()) { + ++(*GetServiceCounters(Counters, "pqproxy|writeSession")->GetCounter("SessionsCreatedTotal", true)); + } + Become(&TThis::StateFunc); + + Database = NKikimr::NPQ::GetDatabaseFromConfig(AppData(ctx)->PQConfig); + const auto& root = AppData(ctx)->PQConfig.GetRoot(); + SelectSourceIdQuery = GetSourceIdSelectQuery(root); + UpdateSourceIdQuery = GetUpdateIdSelectQuery(root); + ConverterFactory = MakeHolder<NPersQueue::TTopicNamesConverterFactory>( + AppData(ctx)->PQConfig, LocalDC + ); + StartTime = ctx.Now(); +} + + +void TWriteSessionActor::Die(const TActorContext& ctx) { + if (State == ES_DYING) + return; + State = ES_DYING; + if (Writer) + ctx.Send(Writer, new TEvents::TEvPoisonPill()); + + if (PipeToBalancer) + NTabletPipe::CloseClient(ctx, PipeToBalancer); + + if (SessionsActive) { + SessionsActive.Dec(); + BytesInflight.Dec(BytesInflight_); + BytesInflightTotal.Dec(BytesInflightTotal_); + } + + LOG_INFO_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "session cookie: " << Cookie << " sessionId: " << OwnerCookie << " is DEAD"); + + if (!Handler->IsShuttingDown()) + Handler->Finish(); + TActorBootstrapped<TWriteSessionActor>::Die(ctx); +} + +void TWriteSessionActor::CheckFinish(const TActorContext& ctx) { + if (!WritesDone) + return; + if (State != ES_INITED) { + CloseSession("out of order Writes done before initialization", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + if (Writes.empty() && FormedWrites.empty() && SentMessages.empty()) { + CloseSession("", NPersQueue::NErrorCode::OK, ctx); + return; + } +} + +void TWriteSessionActor::Handle(TEvPQProxy::TEvDone::TPtr&, const TActorContext& ctx) { + WritesDone = true; + CheckFinish(ctx); +} + +void TWriteSessionActor::CheckACL(const TActorContext& ctx) { + Y_VERIFY(ACLCheckInProgress); + Y_VERIFY(SecurityObject); + NACLib::EAccessRights rights = NACLib::EAccessRights::UpdateRow; + if (!AppData(ctx)->PQConfig.GetCheckACL() || SecurityObject->CheckAccess(rights, *Token)) { + ACLCheckInProgress = false; + if (FirstACLCheck) { + FirstACLCheck = false; + DiscoverPartition(ctx); + } + } else { + TString errorReason = Sprintf("access to topic '%s' denied for '%s' due to 'no WriteTopic rights', Marker# PQ1125", + DiscoveryConverter->GetPrintableString().c_str(), + Token->GetUserSID().c_str()); + CloseSession(errorReason, NPersQueue::NErrorCode::ACCESS_DENIED, ctx); + } +} + +void TWriteSessionActor::Handle(TEvPQProxy::TEvWriteInit::TPtr& ev, const TActorContext& ctx) { + THolder<TEvPQProxy::TEvWriteInit> event(ev->Release()); + + if (State != ES_CREATED) { + //answer error + CloseSession("got second init request", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + const auto& init = event->Request.GetInit(); + + if (init.GetTopic().empty() || init.GetSourceId().empty()) { + CloseSession("no topic or SourceId in init request", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + if (init.GetProxyCookie() != ctx.SelfID.NodeId() && init.GetProxyCookie() != MAGIC_COOKIE_VALUE) { + CloseSession("you must perform ChooseProxy request at first and go to ProxyName server with ProxyCookie", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + //1. Database - !(Root or empty) (Need to bring root DB(s) list to PQConfig) - ONLY search modern path /Database/Path + //2. No database. Try parse and resolve account to database. If possible, try search this path. + //3. Fallback from 2 - legacy mode. + + DiscoveryConverter = ConverterFactory->MakeDiscoveryConverter(init.GetTopic(), true, LocalDC, Database); + if (!DiscoveryConverter->IsValid()) { + CloseSession( + TStringBuilder() << "incorrect topic \"" << DiscoveryConverter->GetOriginalTopic() + << "\": " << DiscoveryConverter->GetReason(), + NPersQueue::NErrorCode::BAD_REQUEST, + ctx + ); + } + PeerName = event->PeerName; + if (!event->Database.empty()) { + Database = event->Database; + } + + SourceId = init.GetSourceId(); + //TODO: check that sourceId does not have characters '"\_% - espace them on client may be? + + Auth = event->Request.GetCredentials(); + event->Request.ClearCredentials(); + Y_PROTOBUF_SUPPRESS_NODISCARD Auth.SerializeToString(&AuthStr); + + LOG_INFO_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "session request cookie: " << Cookie << " " << init << " from " << PeerName); + UserAgent = init.GetVersion(); + LogSession(ctx); + + auto* request = new TEvDescribeTopicsRequest({DiscoveryConverter}); + //TODO: GetNode for /Root/PQ then describe from balancer + ctx.Send(SchemeCache, request); + State = ES_WAIT_SCHEME_2; + InitRequest = init; + PreferedPartition = init.GetPartitionGroup() > 0 ? init.GetPartitionGroup() - 1 : Max<ui32>(); +} + +void TWriteSessionActor::InitAfterDiscovery(const TActorContext& ctx) { + try { + EncodedSourceId = NSourceIdEncoding::EncodeSrcId(FullConverter->GetTopicForSrcIdHash(), SourceId); + } catch (yexception& e) { + CloseSession(TStringBuilder() << "incorrect sourceId \"" << SourceId << "\": " << e.what(), NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + InitMeta = GetInitialDataChunk(InitRequest, FullConverter->GetClientsideName(), PeerName); + + auto subGroup = GetServiceCounters(Counters, "pqproxy|SLI"); + Aggr = {{{{"Account", FullConverter->GetAccount()}}, {"total"}}}; + + SLIErrors = NKikimr::NPQ::TMultiCounter(subGroup, Aggr, {}, {"RequestsError"}, true, "sensor", false); + SLITotal = NKikimr::NPQ::TMultiCounter(subGroup, Aggr, {}, {"RequestsTotal"}, true, "sensor", false); + SLITotal.Inc(); +} + + +void TWriteSessionActor::SetupCounters() +{ + if (SessionsCreated) { + return; + } + + //now topic is checked, can create group for real topic, not garbage + auto subGroup = GetServiceCounters(Counters, "pqproxy|writeSession"); + Y_VERIFY(FullConverter); + auto aggr = GetLabels(FullConverter); + + BytesInflight = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"BytesInflight"}, false); + SessionsWithoutAuth = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"WithoutAuth"}, true); + BytesInflightTotal = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"BytesInflightTotal"}, false); + SessionsCreated = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"SessionsCreated"}, true); + SessionsActive = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"SessionsActive"}, false); + Errors = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"Errors"}, true); + + SessionsCreated.Inc(); + SessionsActive.Inc(); +} + + +void TWriteSessionActor::SetupCounters(const TString& cloudId, const TString& dbId, const TString& folderId) +{ + if (SessionsCreated) { + return; + } + + //now topic is checked, can create group for real topic, not garbage + auto subGroup = GetCountersForStream(Counters); + Y_VERIFY(FullConverter); + auto aggr = GetLabelsForStream(FullConverter, cloudId, dbId, folderId); + + BytesInflight = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"stream.internal_write.bytes_proceeding"}, false, "name"); + SessionsWithoutAuth = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"stream.internal_write.sessions_without_auth"}, true, "name"); + BytesInflightTotal = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"stream.internal_write.bytes_proceeding_total"}, false, "name"); + SessionsCreated = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"stream.internal_write.sessions_created_per_second"}, true, "name"); + SessionsActive = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"stream.internal_write.sessions_active"}, false, "name"); + Errors = NKikimr::NPQ::TMultiCounter(subGroup, aggr, {}, {"stream.internal_write.errors_per_second"}, true, "name"); + + SessionsCreated.Inc(); + SessionsActive.Inc(); +} + + +void TWriteSessionActor::Handle(TEvDescribeTopicsResponse::TPtr& ev, const TActorContext& ctx) { + Y_VERIFY(State == ES_WAIT_SCHEME_2); + auto& res = ev->Get()->Result; + Y_VERIFY(res->ResultSet.size() == 1); + + auto& entry = res->ResultSet[0]; + TString errorReason; + + auto& path = entry.Path; + auto& topic = ev->Get()->TopicsRequested[0]; + switch (entry.Status) { + case TSchemeCacheNavigate::EStatus::RootUnknown: { + errorReason = Sprintf("path '%s' has incorrect root prefix, Marker# PQ14", JoinPath(path).c_str()); + CloseSession(errorReason, NPersQueue::NErrorCode::UNKNOWN_TOPIC, ctx); + return; + } + case TSchemeCacheNavigate::EStatus::PathErrorUnknown: { + errorReason = Sprintf("no path '%s', Marker# PQ151", JoinPath(path).c_str()); + CloseSession(errorReason, NPersQueue::NErrorCode::UNKNOWN_TOPIC, ctx); + return; + } + case TSchemeCacheNavigate::EStatus::Ok: + break; + default: { + errorReason = Sprintf("topic '%s' describe error, Status# %s, Marker# PQ1", path.back().c_str(), + ToString(entry.Status).c_str()); + CloseSession(errorReason, NPersQueue::NErrorCode::ERROR, ctx); + break; + } + } + if (!entry.PQGroupInfo) { + + errorReason = Sprintf("topic '%s' describe error, reason - could not retrieve topic metadata, Marker# PQ99", + topic->GetPrintableString().c_str()); + CloseSession(errorReason, NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + PQInfo = entry.PQGroupInfo; + const auto& description = PQInfo->Description; + //const TString topicName = description.GetName(); + + if (entry.Kind != TSchemeCacheNavigate::EKind::KindTopic) { + errorReason = Sprintf("item '%s' is not a topic, Marker# PQ13", DiscoveryConverter->GetPrintableString().c_str()); + CloseSession(errorReason, NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + FullConverter = DiscoveryConverter->UpgradeToFullConverter(description.GetPQTabletConfig()); + InitAfterDiscovery(ctx); + SecurityObject = entry.SecurityObject; + + Y_VERIFY(description.PartitionsSize() > 0); + + for (ui32 i = 0; i < description.PartitionsSize(); ++i) { + const auto& pi = description.GetPartitions(i); + PartitionToTablet[pi.GetPartitionId()] = pi.GetTabletId(); + } + BalancerTabletId = description.GetBalancerTabletID(); + DatabaseId = description.GetPQTabletConfig().GetYdbDatabaseId(); + FolderId = description.GetPQTabletConfig().GetYcFolderId(); + + if (AppData(ctx)->PQConfig.GetTopicsAreFirstClassCitizen()) { + const auto& tabletConfig = description.GetPQTabletConfig(); + SetupCounters(tabletConfig.GetYcCloudId(), tabletConfig.GetYdbDatabaseId(), + tabletConfig.GetYcFolderId()); + } else { + SetupCounters(); + } + + if (!PipeToBalancer) { + NTabletPipe::TClientConfig clientConfig; + clientConfig.RetryPolicy = { + .RetryLimitCount = 6, + .MinRetryTime = TDuration::MilliSeconds(10), + .MaxRetryTime = TDuration::MilliSeconds(100), + .BackoffMultiplier = 2, + .DoFirstRetryInstantly = true + }; + PipeToBalancer = ctx.RegisterWithSameMailbox(NTabletPipe::CreateClient(ctx.SelfID, BalancerTabletId, clientConfig)); + } + + if (Auth.GetCredentialsCase() == NPersQueueCommon::TCredentials::CREDENTIALS_NOT_SET) { + //ACLCheckInProgress is still true - no recheck will be done + LOG_WARN_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "session without AuthInfo : " << DiscoveryConverter->GetPrintableString() + << " sourceId " << SourceId << " from " << PeerName); + SessionsWithoutAuth.Inc(); + if (AppData(ctx)->PQConfig.GetRequireCredentialsInNewProtocol()) { + CloseSession("Unauthenticated access is forbidden, please provide credentials", NPersQueue::NErrorCode::ACCESS_DENIED, ctx); + return; + } + if (FirstACLCheck) { + FirstACLCheck = false; + DiscoverPartition(ctx); + return; + } + } + + InitCheckACL(ctx); +} + +void TWriteSessionActor::InitCheckACL(const TActorContext& ctx) { + + Y_VERIFY(ACLCheckInProgress); + + TString ticket; + switch (Auth.GetCredentialsCase()) { + case NPersQueueCommon::TCredentials::kTvmServiceTicket: + ticket = Auth.GetTvmServiceTicket(); + break; + case NPersQueueCommon::TCredentials::kOauthToken: + ticket = Auth.GetOauthToken(); + break; + default: + CloseSession("Uknown Credentials case", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + auto entries = GetTicketParserEntries(DatabaseId, FolderId); + ctx.Send(MakeTicketParserID(), new TEvTicketParser::TEvAuthorizeTicket({ + .Database = Database, + .Ticket = ticket, + .PeerName = PeerName, + .Entries = entries + })); +} + +void TWriteSessionActor::Handle(TEvTicketParser::TEvAuthorizeTicketResult::TPtr& ev, const TActorContext& ctx) { + Y_VERIFY(ACLCheckInProgress); + TString ticket = ev->Get()->Ticket; + TString maskedTicket = ticket.size() > 5 ? (ticket.substr(0, 5) + "***" + ticket.substr(ticket.size() - 5)) : "***"; + LOG_INFO_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "CheckACL ticket " << maskedTicket << " got result from TICKET_PARSER response: error: " << ev->Get()->Error << " user: " + << (ev->Get()->Error.empty() ? ev->Get()->Token->GetUserSID() : "")); + + if (!ev->Get()->Error.empty()) { + CloseSession(TStringBuilder() << "Ticket parsing error: " << ev->Get()->Error, NPersQueue::NErrorCode::ACCESS_DENIED, ctx); + return; + } + Token = ev->Get()->Token; + + + Y_VERIFY(ACLCheckInProgress); + LOG_INFO_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "session cookie: " << Cookie << " sessionId: " << OwnerCookie << " describe result for acl check"); + CheckACL(ctx); +} + +void TWriteSessionActor::DiscoverPartition(const NActors::TActorContext& ctx) { + Y_VERIFY(FullConverter); + if (AppData(ctx)->PQConfig.GetTopicsAreFirstClassCitizen()) { + auto partitionId = PreferedPartition; + if (PreferedPartition == Max<ui32>()) { + partitionId = NKikimr::NDataStreams::V1::ShardFromDecimal( + NKikimr::NDataStreams::V1::HexBytesToDecimal(MD5::Calc(SourceId)), PartitionToTablet.size() + ); + } + ProceedPartition(partitionId, ctx); + return; + } + //read from DS + // Hash was always valid here, so new and old are the same + //currently, Topic contains full primary path + SendSelectPartitionRequest(EncodedSourceId.Hash, FullConverter->GetPrimaryPath(), ctx); + //previously topic was like "rt3.dc--account--topic" + SendSelectPartitionRequest(EncodedSourceId.Hash, FullConverter->GetTopicForSrcId(), ctx); + State = ES_WAIT_TABLE_REQUEST_1; +} + +void TWriteSessionActor::SendSelectPartitionRequest(ui32 hash, const TString &topic, + const NActors::TActorContext &ctx +) { + auto ev = MakeHolder<NKqp::TEvKqp::TEvQueryRequest>(); + ev->Record.MutableRequest()->SetAction(NKikimrKqp::QUERY_ACTION_EXECUTE); + ev->Record.MutableRequest()->SetType(NKikimrKqp::QUERY_TYPE_SQL_DML); + ev->Record.MutableRequest()->SetKeepSession(false); + ev->Record.MutableRequest()->SetQuery(SelectSourceIdQuery); + ev->Record.MutableRequest()->SetDatabase(Database); + ev->Record.MutableRequest()->MutableTxControl()->set_commit_tx(true); + ev->Record.MutableRequest()->MutableTxControl()->mutable_begin_tx()->mutable_serializable_read_write(); + ev->Record.MutableRequest()->MutableQueryCachePolicy()->set_keep_in_cache(true); + NClient::TParameters parameters; + parameters["$Hash"] = hash; // 'Valid' hash - short legacy name (account--topic) + parameters["$Topic"] = topic; //currently, Topic contains full primary path + parameters["$SourceId"] = EncodedSourceId.EscapedSourceId; + + ev->Record.MutableRequest()->MutableParameters()->Swap(¶meters); + ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release()); + SelectReqsInflight++; +} + + +void TWriteSessionActor::UpdatePartition(const TActorContext& ctx) { + Y_VERIFY(State == ES_WAIT_TABLE_REQUEST_1 || State == ES_WAIT_NEXT_PARTITION); + auto ev = MakeUpdateSourceIdMetadataRequest(EncodedSourceId.Hash, FullConverter->GetPrimaryPath()); // Now Topic is a path + ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release()); + SourceIdUpdatesInflight++; + + //Previously Topic contained legacy name with DC (rt3.dc1--acc--topic) + ev = MakeUpdateSourceIdMetadataRequest(EncodedSourceId.Hash, FullConverter->GetTopicForSrcId()); + ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release()); + SourceIdUpdatesInflight++; + + State = ES_WAIT_TABLE_REQUEST_2; +} + +void TWriteSessionActor::RequestNextPartition(const TActorContext& ctx) { + Y_VERIFY(State == ES_WAIT_TABLE_REQUEST_1); + State = ES_WAIT_NEXT_PARTITION; + THolder<TEvPersQueue::TEvGetPartitionIdForWrite> x(new TEvPersQueue::TEvGetPartitionIdForWrite); + Y_VERIFY(PipeToBalancer); + + NTabletPipe::SendData(ctx, PipeToBalancer, x.Release()); +} + +void TWriteSessionActor::Handle(TEvPersQueue::TEvGetPartitionIdForWriteResponse::TPtr& ev, const TActorContext& ctx) { + Y_VERIFY(State == ES_WAIT_NEXT_PARTITION); + Partition = ev->Get()->Record.GetPartitionId(); + UpdatePartition(ctx); +} + +void TWriteSessionActor::Handle(NKqp::TEvKqp::TEvQueryResponse::TPtr &ev, const TActorContext &ctx) { + auto& record = ev->Get()->Record.GetRef(); + + if (record.GetYdbStatus() == Ydb::StatusIds::ABORTED) { + LOG_INFO_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "session cookie: " << Cookie << " sessionId: " << OwnerCookie << " sourceID " + << SourceId << " escaped " << EncodedSourceId.EscapedSourceId << " discover partition race, retrying"); + DiscoverPartition(ctx); + return; + } + + if (record.GetYdbStatus() != Ydb::StatusIds::SUCCESS) { + TStringBuilder errorReason; + errorReason << "internal error in kqp Marker# PQ50 : " << record; + if (State == EState::ES_INITED) { + LOG_WARN_S(ctx, NKikimrServices::PQ_WRITE_PROXY, errorReason); + SourceIdUpdatesInflight--; + } else { + CloseSession(errorReason, NPersQueue::NErrorCode::ERROR, ctx); + } + return; + } + + if (State == EState::ES_WAIT_TABLE_REQUEST_1) { + SelectReqsInflight--; + auto& t = record.GetResponse().GetResults(0).GetValue().GetStruct(0); + + if (t.ListSize() != 0) { + auto& tt = t.GetList(0).GetStruct(0); + if (tt.HasOptional() && tt.GetOptional().HasUint32()) { //already got partition + auto accessTime = t.GetList(0).GetStruct(2).GetOptional().GetUint64(); + if (accessTime > MaxSrcIdAccessTime) { // AccessTime + Partition = tt.GetOptional().GetUint32(); + PartitionFound = true; + SourceIdCreateTime = t.GetList(0).GetStruct(1).GetOptional().GetUint64(); + MaxSrcIdAccessTime = accessTime; + } + } + } + if (SelectReqsInflight != 0) { + return; + } + if (SourceIdCreateTime == 0) { + SourceIdCreateTime = TInstant::Now().MilliSeconds(); + } + if (PartitionFound && PreferedPartition < Max<ui32>() && Partition != PreferedPartition) { + CloseSession(TStringBuilder() << "SourceId " << SourceId << " is already bound to PartitionGroup " << (Partition + 1) << ", but client provided " << (PreferedPartition + 1) << ". SourceId->PartitionGroup binding cannot be changed, either use another SourceId, specify PartitionGroup " << (Partition + 1) << ", or do not specify PartitionGroup at all.", + NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + LOG_INFO_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "session cookie: " << Cookie << " sessionId: " << OwnerCookie << " sourceID " + << SourceId << " escaped " << EncodedSourceId.EscapedSourceId << " hash " << EncodedSourceId.Hash << " partition " << Partition << " partitions " + << PartitionToTablet.size() << "(" << EncodedSourceId.Hash % PartitionToTablet.size() << ") create " << SourceIdCreateTime << " result " << t); + + if (!PartitionFound && (PreferedPartition < Max<ui32>() || !AppData(ctx)->PQConfig.GetRoundRobinPartitionMapping())) { + Partition = PreferedPartition < Max<ui32>() ? PreferedPartition : EncodedSourceId.Hash % PartitionToTablet.size(); //choose partition default value + PartitionFound = true; + } + + if (PartitionFound) { + UpdatePartition(ctx); + } else { + RequestNextPartition(ctx); + } + return; + } else if (State == EState::ES_WAIT_TABLE_REQUEST_2) { + Y_VERIFY(SourceIdUpdatesInflight > 0); + SourceIdUpdatesInflight--; + if (SourceIdUpdatesInflight == 0) { + LastSourceIdUpdate = ctx.Now(); + ProceedPartition(Partition, ctx); + } + } else if (State == EState::ES_INITED) { + Y_VERIFY(SourceIdUpdatesInflight > 0); + SourceIdUpdatesInflight--; + if (SourceIdUpdatesInflight == 0) { + LastSourceIdUpdate = ctx.Now(); + } + } else { + Y_FAIL("Wrong state"); + } +} + +THolder<NKqp::TEvKqp::TEvQueryRequest> TWriteSessionActor::MakeUpdateSourceIdMetadataRequest( + ui32 hash, const TString& topic +) { + auto ev = MakeHolder<NKqp::TEvKqp::TEvQueryRequest>(); + + ev->Record.MutableRequest()->SetAction(NKikimrKqp::QUERY_ACTION_EXECUTE); + ev->Record.MutableRequest()->SetType(NKikimrKqp::QUERY_TYPE_SQL_DML); + ev->Record.MutableRequest()->SetQuery(UpdateSourceIdQuery); + ev->Record.MutableRequest()->SetDatabase(Database); + ev->Record.MutableRequest()->SetKeepSession(false); + ev->Record.MutableRequest()->MutableTxControl()->set_commit_tx(true); + ev->Record.MutableRequest()->MutableTxControl()->mutable_begin_tx()->mutable_serializable_read_write(); + ev->Record.MutableRequest()->MutableQueryCachePolicy()->set_keep_in_cache(true); + + NClient::TParameters parameters; + parameters["$Hash"] = hash; + parameters["$Topic"] = topic; //Previously Topic contained legacy name with DC (rt3.dc1--acc--topic) + parameters["$SourceId"] = EncodedSourceId.EscapedSourceId; + parameters["$CreateTime"] = SourceIdCreateTime; + parameters["$AccessTime"] = TInstant::Now().MilliSeconds(); + parameters["$Partition"] = Partition; + ev->Record.MutableRequest()->MutableParameters()->Swap(¶meters); + + return ev; +} + + +void TWriteSessionActor::Handle(NKqp::TEvKqp::TEvProcessResponse::TPtr &ev, const TActorContext &ctx) { + auto& record = ev->Get()->Record; + LOG_INFO_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "session cookie: " << Cookie << " sessionId: " << OwnerCookie << " sourceID " + << SourceId << " escaped " << EncodedSourceId.EscapedSourceId << " discover partition error - " << record); + CloseSession("Internal error on discovering partition", NPersQueue::NErrorCode::ERROR, ctx); +} + + +void TWriteSessionActor::ProceedPartition(const ui32 partition, const TActorContext& ctx) { + Partition = partition; + auto it = PartitionToTablet.find(Partition); + + ui64 tabletId = it != PartitionToTablet.end() ? it->second : 0; + + if (!tabletId) { + CloseSession( + Sprintf("no partition %u in topic '%s', Marker# PQ4", Partition, DiscoveryConverter->GetPrintableString().c_str()), + NPersQueue::NErrorCode::UNKNOWN_TOPIC, ctx + ); + return; + } + + Writer = ctx.RegisterWithSameMailbox(NPQ::CreatePartitionWriter(ctx.SelfID, tabletId, Partition, SourceId)); + State = ES_WAIT_WRITER_INIT; + + ui32 border = AppData(ctx)->PQConfig.GetWriteInitLatencyBigMs(); + auto subGroup = GetServiceCounters(Counters, "pqproxy|SLI"); + + InitLatency = NKikimr::NPQ::CreateSLIDurationCounter(subGroup, Aggr, "WriteInit", border, {100, 200, 500, 1000, 1500, 2000, 5000, 10000, 30000, 99999999}); + SLIBigLatency = NKikimr::NPQ::TMultiCounter(subGroup, Aggr, {}, {"RequestsBigLatency"}, true, "sensor", false); + + ui32 initDurationMs = (ctx.Now() - StartTime).MilliSeconds(); + InitLatency.IncFor(initDurationMs, 1); + if (initDurationMs >= border) { + SLIBigLatency.Inc(); + } +} + +void TWriteSessionActor::CloseSession(const TString& errorReason, const NPersQueue::NErrorCode::EErrorCode errorCode, const NActors::TActorContext& ctx) { + if (errorCode != NPersQueue::NErrorCode::OK) { + if (InternalErrorCode(errorCode)) { + SLIErrors.Inc(); + } + + if (Errors) { + Errors.Inc(); + } else if (!AppData(ctx)->PQConfig.GetTopicsAreFirstClassCitizen()) { + ++(*GetServiceCounters(Counters, "pqproxy|writeSession")->GetCounter("Errors", true)); + } + + TWriteResponse result; + + auto error = result.MutableError(); + error->SetDescription(errorReason); + error->SetCode(errorCode); + + LOG_INFO_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "session error cookie: " << Cookie << " reason: \"" << errorReason << "\" code: " << EErrorCode_Name(errorCode) << " sessionId: " << OwnerCookie); + + Handler->Reply(result); + } else { + LOG_INFO_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "session closed cookie: " << Cookie << " sessionId: " << OwnerCookie); + } + + Die(ctx); +} + +void TWriteSessionActor::Handle(NPQ::TEvPartitionWriter::TEvInitResult::TPtr& ev, const TActorContext& ctx) { + if (State != ES_WAIT_WRITER_INIT) { + return CloseSession("got init result but not wait for it", NPersQueue::NErrorCode::ERROR, ctx); + } + + const auto& result = *ev->Get(); + if (!result.IsSuccess()) { + const auto& error = result.GetError(); + if (error.Response.HasErrorCode()) { + return CloseSession("status is not ok: " + error.Response.GetErrorReason(), error.Response.GetErrorCode(), ctx); + } else { + return CloseSession("error at writer init: " + error.Reason, NPersQueue::NErrorCode::ERROR, ctx); + } + } + + OwnerCookie = result.GetResult().OwnerCookie; + const auto& maxSeqNo = result.GetResult().SourceIdInfo.GetSeqNo(); + + TWriteResponse response; + auto init = response.MutableInit(); + init->SetSessionId(EscapeC(OwnerCookie)); + init->SetMaxSeqNo(maxSeqNo); + init->SetPartition(Partition); + Y_VERIFY(FullConverter); + init->SetTopic(FullConverter->GetClientsideName()); + + LOG_INFO_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "session inited cookie: " << Cookie << " partition: " << Partition + << " MaxSeqNo: " << maxSeqNo << " sessionId: " << OwnerCookie); + + Handler->Reply(response); + + State = ES_INITED; + + ctx.Schedule(CHECK_ACL_DELAY, new TEvents::TEvWakeup()); + + //init completed; wait for first data chunk + NextRequestInited = true; + Handler->ReadyForNextRead(); +} + +void TWriteSessionActor::Handle(NPQ::TEvPartitionWriter::TEvWriteAccepted::TPtr& ev, const TActorContext& ctx) { + if (State != ES_INITED) { + return CloseSession("got write permission but not wait for it", NPersQueue::NErrorCode::ERROR, ctx); + } + + Y_VERIFY(!FormedWrites.empty()); + TWriteRequestBatchInfo::TPtr writeRequest = std::move(FormedWrites.front()); + + if (ev->Get()->Cookie != writeRequest->Cookie) { + return CloseSession("out of order reserve bytes response from server, may be previous is lost", NPersQueue::NErrorCode::ERROR, ctx); + } + + FormedWrites.pop_front(); + + ui64 diff = writeRequest->ByteSize; + + SentMessages.emplace_back(std::move(writeRequest)); + + BytesInflight_ -= diff; + BytesInflight.Dec(diff); + + if (!NextRequestInited && BytesInflight_ < MAX_BYTES_INFLIGHT) { //allow only one big request to be readed but not sended + NextRequestInited = true; + Handler->ReadyForNextRead(); + } + + --NumReserveBytesRequests; + if (!Writes.empty()) + GenerateNextWriteRequest(ctx); +} + +void TWriteSessionActor::Handle(NPQ::TEvPartitionWriter::TEvWriteResponse::TPtr& ev, const TActorContext& ctx) { + if (State != ES_INITED) { + return CloseSession("got write response but not wait for it", NPersQueue::NErrorCode::ERROR, ctx); + } + + const auto& result = *ev->Get(); + if (!result.IsSuccess()) { + const auto& record = result.Record; + if (record.HasErrorCode()) { + return CloseSession("status is not ok: " + record.GetErrorReason(), record.GetErrorCode(), ctx); + } else { + return CloseSession("error at write: " + result.GetError().Reason, NPersQueue::NErrorCode::ERROR, ctx); + } + } + + const auto& resp = result.Record.GetPartitionResponse(); + + if (SentMessages.empty()) { + CloseSession("got too many replies from server, internal error", NPersQueue::NErrorCode::ERROR, ctx); + return; + } + + TWriteRequestBatchInfo::TPtr writeRequest = std::move(SentMessages.front()); + SentMessages.pop_front(); + + if (resp.GetCookie() != writeRequest->Cookie) { + return CloseSession("out of order write response from server, may be previous is lost", NPersQueue::NErrorCode::ERROR, ctx); + } + + auto addAck = [](const TPersQueuePartitionResponse::TCmdWriteResult& res, TWriteResponse::TAck* ack, TWriteResponse::TStat* stat) { + ack->SetSeqNo(res.GetSeqNo()); + ack->SetOffset(res.GetOffset()); + ack->SetAlreadyWritten(res.GetAlreadyWritten()); + + stat->SetTotalTimeInPartitionQueueMs( + Max(res.GetTotalTimeInPartitionQueueMs(), stat->GetTotalTimeInPartitionQueueMs())); + stat->SetPartitionQuotedTimeMs( + Max(res.GetPartitionQuotedTimeMs(), stat->GetPartitionQuotedTimeMs())); + stat->SetTopicQuotedTimeMs( + Max(res.GetTopicQuotedTimeMs(), stat->GetTopicQuotedTimeMs())); + stat->SetWriteTimeMs( + Max(res.GetWriteTimeMs(), stat->GetWriteTimeMs())); + }; + + size_t cmdWriteResultIndex = 0; + for (const auto& userWriteRequest : writeRequest->UserWriteRequests) { + TWriteResponse result; + if (userWriteRequest->Request.HasDataBatch()) { + if (resp.CmdWriteResultSize() - cmdWriteResultIndex < userWriteRequest->Request.GetDataBatch().DataSize()) { + CloseSession("too less responses from server", NPersQueue::NErrorCode::ERROR, ctx); + return; + } + for (size_t endIndex = cmdWriteResultIndex + userWriteRequest->Request.GetDataBatch().DataSize(); cmdWriteResultIndex < endIndex; ++cmdWriteResultIndex) { + addAck(resp.GetCmdWriteResult(cmdWriteResultIndex), + result.MutableAckBatch()->AddAck(), + result.MutableAckBatch()->MutableStat()); + } + } else { + Y_VERIFY(userWriteRequest->Request.HasData()); + if (cmdWriteResultIndex >= resp.CmdWriteResultSize()) { + CloseSession("too less responses from server", NPersQueue::NErrorCode::ERROR, ctx); + return; + } + auto* ack = result.MutableAck(); + addAck(resp.GetCmdWriteResult(cmdWriteResultIndex), ack, ack->MutableStat()); + ++cmdWriteResultIndex; + } + Handler->Reply(result); + } + + ui64 diff = writeRequest->ByteSize; + + BytesInflightTotal_ -= diff; + BytesInflightTotal.Dec(diff); + + CheckFinish(ctx); +} + +void TWriteSessionActor::Handle(NPQ::TEvPartitionWriter::TEvDisconnected::TPtr&, const TActorContext& ctx) { + CloseSession("pipe to partition's tablet is dead", NPersQueue::NErrorCode::ERROR, ctx); +} + +void TWriteSessionActor::Handle(TEvTabletPipe::TEvClientConnected::TPtr& ev, const TActorContext& ctx) { + TEvTabletPipe::TEvClientConnected *msg = ev->Get(); + if (msg->Status != NKikimrProto::OK) { + CloseSession(TStringBuilder() << "pipe to tablet is dead " << msg->TabletId, NPersQueue::NErrorCode::ERROR, ctx); + return; + } +} + +void TWriteSessionActor::Handle(TEvTabletPipe::TEvClientDestroyed::TPtr& ev, const TActorContext& ctx) { + CloseSession(TStringBuilder() << "pipe to tablet is dead " << ev->Get()->TabletId, NPersQueue::NErrorCode::ERROR, ctx); +} + +void TWriteSessionActor::GenerateNextWriteRequest(const TActorContext& ctx) { + TWriteRequestBatchInfo::TPtr writeRequest = new TWriteRequestBatchInfo(); + + auto ev = MakeHolder<NPQ::TEvPartitionWriter::TEvWriteRequest>(++NextRequestCookie); + NKikimrClient::TPersQueueRequest& request = ev->Record; + + writeRequest->UserWriteRequests = std::move(Writes); + Writes.clear(); + + i64 diff = 0; + auto addData = [&](const TWriteRequest::TData& data) { + auto w = request.MutablePartitionRequest()->AddCmdWrite(); + w->SetData(GetSerializedData(InitMeta, data)); + w->SetClientDC(ClientDC); + w->SetSeqNo(data.GetSeqNo()); + w->SetSourceId(NPQ::NSourceIdEncoding::EncodeSimple(SourceId)); // EncodeSimple is needed for compatibility with LB + //TODO: add in SourceID clientId when TVM will be ready + w->SetCreateTimeMS(data.GetCreateTimeMs()); + w->SetUncompressedSize(data.GetUncompressedSize()); + }; + + for (const auto& write : writeRequest->UserWriteRequests) { + diff -= write->Request.ByteSize(); + if (write->Request.HasDataBatch()) { + for (const TWriteRequest::TData& data : write->Request.GetDataBatch().GetData()) { + addData(data); + } + } else { // single data + Y_VERIFY(write->Request.HasData()); + addData(write->Request.GetData()); + } + } + + writeRequest->Cookie = request.GetPartitionRequest().GetCookie(); + + Y_VERIFY(-diff <= (i64)BytesInflight_); + diff += request.ByteSize(); + BytesInflight_ += diff; + BytesInflightTotal_ += diff; + BytesInflight.Inc(diff); + BytesInflightTotal.Inc(diff); + + writeRequest->ByteSize = request.ByteSize(); + FormedWrites.push_back(writeRequest); + + ctx.Send(Writer, std::move(ev)); + ++NumReserveBytesRequests; +} + +TString TWriteSessionActor::CheckSupportedCodec(const ui32 codecId) { + TString err; + const auto& description = PQInfo->Description; + if (!description.GetPQTabletConfig().HasCodecs() || description.GetPQTabletConfig().GetCodecs().IdsSize() == 0) + return ""; + + Y_VERIFY(description.PartitionsSize() > 0); + for (const auto& codec : description.GetPQTabletConfig().GetCodecs().GetIds()) { + if (codecId == codec) { + return ""; + } + } + err = "Unsupported codec provided. Supported codecs for this topic are:"; + bool first = true; + for (const auto& codec : description.GetPQTabletConfig().GetCodecs().GetCodecs()) { + if (first) { + first = false; + } else { + err += ","; + } + err += " " + codec; + } + return err; +} + + +void TWriteSessionActor::Handle(TEvPQProxy::TEvWrite::TPtr& ev, const TActorContext& ctx) { + + RequestNotChecked = true; + + if (State != ES_INITED) { + //answer error + CloseSession("write in not inited session", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return; + } + + auto auth = ev->Get()->Request.GetCredentials(); + ev->Get()->Request.ClearCredentials(); + TString tmp; + Y_PROTOBUF_SUPPRESS_NODISCARD auth.SerializeToString(&tmp); + if (auth.GetCredentialsCase() != NPersQueueCommon::TCredentials::CREDENTIALS_NOT_SET && tmp != AuthStr) { + Auth = auth; + AuthStr = tmp; + ForceACLCheck = true; + } + auto dataCheck = [&](const TWriteRequest::TData& data) -> bool { + if (!data.GetSeqNo()) { + CloseSession("bad write request - SeqNo must be positive", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return false; + } + + if (data.GetData().empty()) { + CloseSession("bad write request - data must be non-empty", NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return false; + } + TString err = CheckSupportedCodec((ui32)data.GetCodec()); + if (!err.empty()) { + CloseSession(err, NPersQueue::NErrorCode::BAD_REQUEST, ctx); + return false; + } + + return true; + }; + if (ev->Get()->Request.HasDataBatch()) { + for (const auto& data : ev->Get()->Request.GetDataBatch().GetData()) { + if (!dataCheck(data)) { + return; + } + } + } else { + Y_VERIFY(ev->Get()->Request.HasData()); + if (!dataCheck(ev->Get()->Request.GetData())) { + return; + } + } + + THolder<TEvPQProxy::TEvWrite> event(ev->Release()); + Writes.push_back(std::move(event)); + + ui64 diff = Writes.back()->Request.ByteSize(); + BytesInflight_ += diff; + BytesInflightTotal_ += diff; + BytesInflight.Inc(diff); + BytesInflightTotal.Inc(diff); + + if (BytesInflight_ < MAX_BYTES_INFLIGHT) { //allow only one big request to be readed but not sended + Y_VERIFY(NextRequestInited); + Handler->ReadyForNextRead(); + } else { + NextRequestInited = false; + } + + if (NumReserveBytesRequests < MAX_RESERVE_REQUESTS_INFLIGHT) { + GenerateNextWriteRequest(ctx); + } +} + + +void TWriteSessionActor::HandlePoison(TEvPQProxy::TEvDieCommand::TPtr& ev, const TActorContext& ctx) { + CloseSession(ev->Get()->Reason, ev->Get()->ErrorCode, ctx); +} + + +void TWriteSessionActor::LogSession(const TActorContext& ctx) { + + LOG_INFO_S(ctx, NKikimrServices::PQ_WRITE_PROXY, "write session: cookie=" << Cookie << " sessionId=" << OwnerCookie + << " userAgent=\"" << UserAgent << "\" ip=" << PeerName << " proto=v0 " + << " topic=" << DiscoveryConverter->GetPrintableString() << " durationSec=" << (ctx.Now() - StartTime).Seconds()); + + LogSessionDeadline = ctx.Now() + TDuration::Hours(1) + TDuration::Seconds(rand() % 60); +} + +void TWriteSessionActor::HandleWakeup(const TActorContext& ctx) { + Y_VERIFY(State == ES_INITED); + ctx.Schedule(CHECK_ACL_DELAY, new TEvents::TEvWakeup()); + if (!ACLCheckInProgress && (ForceACLCheck || (ctx.Now() - LastACLCheckTimestamp > TDuration::Seconds(AppData(ctx)->PQConfig.GetACLRetryTimeoutSec()) && RequestNotChecked))) { + ForceACLCheck = false; + RequestNotChecked = false; + if (Auth.GetCredentialsCase() != NPersQueueCommon::TCredentials::CREDENTIALS_NOT_SET) { + ACLCheckInProgress = true; + auto* request = new TEvDescribeTopicsRequest({DiscoveryConverter}); + ctx.Send(SchemeCache, request); + } + } + if (!SourceIdUpdatesInflight && ctx.Now() - LastSourceIdUpdate > SOURCEID_UPDATE_PERIOD) { + SourceIdUpdatesInflight++; + Y_VERIFY(FullConverter); + auto ev = MakeUpdateSourceIdMetadataRequest(EncodedSourceId.Hash, FullConverter->GetPrimaryPath()); // Now Topic is a path + ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release()); + // Previously Topic contained legacy name with DC (rt3.dc1--acc--topic) + SourceIdUpdatesInflight++; + ev = MakeUpdateSourceIdMetadataRequest(EncodedSourceId.Hash, FullConverter->GetTopicForSrcId()); + ctx.Send(NKqp::MakeKqpProxyID(ctx.SelfID.NodeId()), ev.Release()); + } + if (ctx.Now() >= LogSessionDeadline) { + LogSession(ctx); + } +} + +} +} diff --git a/kikimr/yndx/grpc_services/persqueue/persqueue.cpp b/kikimr/yndx/grpc_services/persqueue/persqueue.cpp new file mode 100644 index 0000000000..43b406a279 --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/persqueue.cpp @@ -0,0 +1,59 @@ +#include "persqueue.h" +#include "grpc_pq_read.h" +#include "grpc_pq_write.h" + +#include <ydb/core/base/appdata.h> +#include <ydb/core/base/counters.h> + + +namespace NKikimr { +namespace NGRpcService { + +static const ui32 PersQueueWriteSessionsMaxCount = 1000000; +static const ui32 PersQueueReadSessionsMaxCount = 100000; + +TGRpcPersQueueService::TGRpcPersQueueService(NActors::TActorSystem *system, TIntrusivePtr<NMonitoring::TDynamicCounters> counters, const NActors::TActorId& schemeCache) + : ActorSystem(system) + , Counters(counters) + , SchemeCache(schemeCache) +{ } + +void TGRpcPersQueueService::InitService(grpc::ServerCompletionQueue *cq, NGrpc::TLoggerPtr) { + CQ = cq; + if (ActorSystem->AppData<TAppData>()->PQConfig.GetEnabled()) { + WriteService.reset(new NGRpcProxy::TPQWriteService(GetService(), CQ, ActorSystem, SchemeCache, Counters, PersQueueWriteSessionsMaxCount)); + WriteService->InitClustersUpdater(); + ReadService.reset(new NGRpcProxy::TPQReadService(this, CQ, ActorSystem, SchemeCache, Counters, PersQueueReadSessionsMaxCount)); + SetupIncomingRequests(); + } +} + +void TGRpcPersQueueService::SetGlobalLimiterHandle(NGrpc::TGlobalLimiter* limiter) { + Limiter = limiter; +} + +bool TGRpcPersQueueService::IncRequest() { + return Limiter->Inc(); +} + +void TGRpcPersQueueService::DecRequest() { + Limiter->Dec(); +} + +void TGRpcPersQueueService::SetupIncomingRequests() { + WriteService->SetupIncomingRequests(); + ReadService->SetupIncomingRequests(); +} + +void TGRpcPersQueueService::StopService() noexcept { + TGrpcServiceBase::StopService(); + if (WriteService.get() != nullptr) { + WriteService->StopService(); + } + if (ReadService.get() != nullptr) { + ReadService->StopService(); + } +} + +} // namespace NGRpcService +} // namespace NKikimr diff --git a/kikimr/yndx/grpc_services/persqueue/persqueue.h b/kikimr/yndx/grpc_services/persqueue/persqueue.h new file mode 100644 index 0000000000..267efa7a6d --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/persqueue.h @@ -0,0 +1,49 @@ +#pragma once + +#include <library/cpp/actors/core/actorsystem.h> + +#include <kikimr/yndx/api/grpc/persqueue.grpc.pb.h> + +#include <library/cpp/grpc/server/grpc_server.h> + + +namespace NKikimr { + +namespace NGRpcProxy { + class TPQWriteService; + class TPQReadService; +} + +namespace NGRpcService { + +class TGRpcPersQueueService + : public NGrpc::TGrpcServiceBase<NPersQueue::PersQueueService> +{ +public: + TGRpcPersQueueService(NActors::TActorSystem* system, TIntrusivePtr<NMonitoring::TDynamicCounters> counters, const NActors::TActorId& schemeCache); + + void InitService(grpc::ServerCompletionQueue* cq, NGrpc::TLoggerPtr logger) override; + void SetGlobalLimiterHandle(NGrpc::TGlobalLimiter* limiter) override; + void StopService() noexcept override; + + using NGrpc::TGrpcServiceBase<NPersQueue::PersQueueService>::GetService; + + bool IncRequest(); + void DecRequest(); + +private: + void SetupIncomingRequests(); + + NActors::TActorSystem* ActorSystem; + grpc::ServerCompletionQueue* CQ = nullptr; + + TIntrusivePtr<NMonitoring::TDynamicCounters> Counters; + NGrpc::TGlobalLimiter* Limiter = nullptr; + NActors::TActorId SchemeCache; + + std::shared_ptr<NGRpcProxy::TPQWriteService> WriteService; + std::shared_ptr<NGRpcProxy::TPQReadService> ReadService; +}; + +} // namespace NGRpcService +} // namespace NKikimr diff --git a/kikimr/yndx/grpc_services/persqueue/persqueue_compat_ut.cpp b/kikimr/yndx/grpc_services/persqueue/persqueue_compat_ut.cpp new file mode 100644 index 0000000000..7b9b117bcd --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/persqueue_compat_ut.cpp @@ -0,0 +1,122 @@ +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_pqlib.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/data_writer.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/types.h> + +namespace NKikimr::NPersQueueTests { + +using namespace NPersQueue::NTests; + +class TPQv0CompatTestBase { +public: + THolder<TTestPQLib> PQLib; + THolder<::NPersQueue::TTestServer> Server; + TString OriginalLegacyName1; + TString OriginalModernName1; + TString MirroredLegacyName1; + TString MirroredModernName1; + TString ShortName1; + + TString OriginalLegacyName2; + TString OriginalModernName2; + TString MirroredLegacyName2; + TString MirroredModernName2; + TString ShortName2; + +public: + TPQv0CompatTestBase() + { + Server = MakeHolder<::NPersQueue::TTestServer>(false); + Server->ServerSettings.PQConfig.MutablePQDiscoveryConfig()->SetLbUserDatabaseRoot("/Root/LB"); + Server->ServerSettings.PQConfig.SetCheckACL(false); + Server->StartServer(); + Server->EnableLogs({ NKikimrServices::KQP_PROXY }, NActors::NLog::PRI_EMERG); + Server->EnableLogs({ NKikimrServices::PERSQUEUE }, NActors::NLog::PRI_INFO); + Server->EnableLogs({ NKikimrServices::PQ_METACACHE }, NActors::NLog::PRI_DEBUG); + OriginalLegacyName1 = "rt3.dc1--account--topic1"; + MirroredLegacyName1 = "rt3.dc2--account--topic1"; + OriginalModernName1 = "/Root/LB/account/topic1"; + MirroredModernName1 = "/Root/LB/account/.topic2/mirrored-from-dc2"; + ShortName1 = "account/topic1"; + + OriginalLegacyName2 = "rt3.dc1--account--topic2"; + MirroredLegacyName2 = "rt3.dc2--account--topic2"; + OriginalModernName2 = "/Root/LB/account/topic2"; + MirroredModernName2 = "/Root/LB/account/.topic2/mirrored-from-dc2"; + ShortName2 = "account/topic2"; + + Server->AnnoyingClient->CreateTopicNoLegacy(OriginalLegacyName1, 1, false); + Server->AnnoyingClient->CreateTopicNoLegacy(MirroredLegacyName1, 1, false); + Server->AnnoyingClient->CreateTopicNoLegacy(OriginalModernName2, 1, true, true, "dc1"); + Server->AnnoyingClient->CreateTopicNoLegacy(MirroredModernName2, 1, true, true, "dc2"); + Server->AnnoyingClient->CreateConsumer("test-consumer"); + InitPQLib(); + } + void InitPQLib() { + PQLib = MakeHolder<TTestPQLib>(*Server); + TPQDataWriter writer{OriginalLegacyName1, ShortName1, "test", *Server}; + writer.WaitWritePQServiceInitialization(); + }; +}; + +Y_UNIT_TEST_SUITE(TPQCompatTest) { + Y_UNIT_TEST(DiscoverTopics) { + TPQv0CompatTestBase testServer; + Cerr << "Create producer\n"; + { + auto [producer, res] = testServer.PQLib->CreateProducer(testServer.ShortName2, "123", {}, ::NPersQueue::ECodec::RAW); + Cerr << "Got response: " << res.Response.ShortDebugString() << Endl; + UNIT_ASSERT(res.Response.HasInit()); + } + Cerr << "Create producer(2)\n"; + { + auto [producer, res] = testServer.PQLib->CreateProducer(testServer.ShortName1, "123", {}, ::NPersQueue::ECodec::RAW); + UNIT_ASSERT(res.Response.HasInit()); + } + } + + Y_UNIT_TEST(SetupLockSession) { + TPQv0CompatTestBase server{}; + auto [consumer, startResult] = server.PQLib->CreateConsumer({server.ShortName1}, "test-consumer", 1, true); + Cerr << startResult.Response << "\n"; + for (ui32 i = 0; i < 2; ++i) { + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << "Response: " << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetTopic() == server.OriginalLegacyName1 + || msg.GetValue().Response.GetLock().GetTopic() == server.MirroredLegacyName1); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetPartition() == 0); + } + auto msg = consumer->GetNextMessage(); + UNIT_ASSERT(!msg.Wait(TDuration::Seconds(1))); + server.Server->AnnoyingClient->AlterTopic(server.MirroredLegacyName1, 2); + msg.Wait(); + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetTopic() == server.MirroredLegacyName1); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetPartition() == 1); + } + + Y_UNIT_TEST(LegacyRequests) { + TPQv0CompatTestBase server{}; + server.Server->AnnoyingClient->GetPartOffset( + { + {server.OriginalLegacyName1, {0}}, + {server.MirroredLegacyName1, {0}}, + {server.OriginalLegacyName2, {0}}, + {server.MirroredLegacyName2, {0}}, + }, + 4, 0, true + ); + server.Server->AnnoyingClient->SetClientOffsetPQ(server.OriginalLegacyName2, 0, 5); + server.Server->AnnoyingClient->SetClientOffsetPQ(server.MirroredLegacyName2, 0, 5); + + server.Server->AnnoyingClient->GetPartOffset( + { + {server.OriginalLegacyName2, {0}}, + {server.MirroredLegacyName2, {0}}, + }, + 2, 2, true + ); + } +} +} //namespace NKikimr::NPersQueueTests; diff --git a/kikimr/yndx/grpc_services/persqueue/persqueue_ut.cpp b/kikimr/yndx/grpc_services/persqueue/persqueue_ut.cpp new file mode 100644 index 0000000000..6a899a684d --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/persqueue_ut.cpp @@ -0,0 +1,2405 @@ +#include "ut/definitions.h" +#include <ydb/core/base/appdata.h> +#include <ydb/core/testlib/test_pq_client.h> +#include <ydb/core/client/server/grpc_proxy_status.h> +#include <ydb/core/protos/grpc_pq_old.pb.h> +#include <ydb/core/persqueue/writer/source_id_encoding.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_pqlib.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h> + +#include <ydb/library/aclib/aclib.h> +#include <ydb/library/persqueue/tests/counters.h> + +#include <library/cpp/testing/unittest/tests_data.h> +#include <library/cpp/testing/unittest/registar.h> +#include <library/cpp/http/io/stream.h> +#include <google/protobuf/text_format.h> + +#include <util/string/join.h> +#include <util/string/builder.h> + +#include <grpc++/client_context.h> +#include <grpc++/create_channel.h> + +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h> + +namespace NKikimr { +namespace NPersQueueTests { + +using namespace Tests; +using namespace NKikimrClient; +using namespace NPersQueue; +using namespace NPersQueue::NTests; +using namespace NThreading; +using namespace NNetClassifier; + +static TString FormNetData() { + return "10.99.99.224/32\tSAS\n" + "::1/128\tVLA\n"; +} + +TAutoPtr<IEventHandle> GetClassifierUpdate(TServer& server, const TActorId sender) { + auto& actorSystem = *server.GetRuntime(); + actorSystem.Send( + new IEventHandle(MakeNetClassifierID(), sender, + new TEvNetClassifier::TEvSubscribe() + )); + + TAutoPtr<IEventHandle> handle; + actorSystem.GrabEdgeEvent<NNetClassifier::TEvNetClassifier::TEvClassifierUpdate>(handle); + + UNIT_ASSERT(handle); + UNIT_ASSERT_VALUES_EQUAL(handle->Recipient, sender); + + return handle; +} + +THolder<TTempFileHandle> CreateNetDataFile(const TString& content) { + auto netDataFile = MakeHolder<TTempFileHandle>("data.tsv"); + + netDataFile->Write(content.Data(), content.Size()); + netDataFile->FlushData(); + + return netDataFile; +} + + +Y_UNIT_TEST_SUITE(TPersQueueTest2) { + void PrepareForGrpcNoDC(TFlatMsgBusPQClient& annoyingClient) { + annoyingClient.SetNoConfigMode(); + annoyingClient.FullInit(); + annoyingClient.InitUserRegistry(); + annoyingClient.MkDir("/Root", "account1"); + annoyingClient.MkDir("/Root/PQ", "account1"); + annoyingClient.CreateTopicNoLegacy("/Root/PQ/rt3.db--topic1", 5, false); + annoyingClient.CreateTopicNoLegacy("/Root/PQ/account1/topic1", 5, false, true, Nothing(), {"user1", "user2"}); + annoyingClient.CreateTopicNoLegacy("/Root/account2/topic2", 5); + } + Y_UNIT_TEST(TestGrpcWriteNoDC) { + TTestServer server(false); + server.ServerSettings.PQConfig.SetTopicsAreFirstClassCitizen(true); + server.ServerSettings.PQConfig.SetRoot("/Rt2/PQ"); + server.ServerSettings.PQConfig.SetDatabase("/Root"); + server.StartServer(); + + UNIT_ASSERT_VALUES_EQUAL(NMsgBusProxy::MSTATUS_OK, + server.AnnoyingClient->AlterUserAttributes("/", "Root", {{"folder_id", "somefolder"}, {"cloud_id", "somecloud"}, {"database_id", "root"}})); + + + PrepareForGrpcNoDC(*server.AnnoyingClient); + auto writer = MakeDataWriter(server, "source1"); + + writer.Write("/Root/account2/topic2", {"valuevaluevalue1"}, true, "topic1@" BUILTIN_ACL_DOMAIN); + writer.Write("/Root/PQ/account1/topic1", {"valuevaluevalue1"}, true, "topic1@" BUILTIN_ACL_DOMAIN); + + NACLib::TDiffACL acl; + acl.AddAccess(NACLib::EAccessType::Allow, NACLib::UpdateRow, "topic1@" BUILTIN_ACL_DOMAIN); + acl.AddAccess(NACLib::EAccessType::Allow, NACLib::SelectRow, "topic1@" BUILTIN_ACL_DOMAIN); + + server.AnnoyingClient->ModifyACL("/Root/account2", "topic2", acl.SerializeAsString()); + server.AnnoyingClient->ModifyACL("/Root/PQ/account1", "topic1", acl.SerializeAsString()); + + Sleep(TDuration::Seconds(5)); + + writer.Write("/Root/account2/topic2", {"valuevaluevalue1"}, false, "topic1@" BUILTIN_ACL_DOMAIN); + + writer.Write("/Root/PQ/account1/topic1", {"valuevaluevalue1"}, false, "topic1@" BUILTIN_ACL_DOMAIN); + writer.Write("/Root/PQ/account1/topic1", {"valuevaluevalue2"}, false, "topic1@" BUILTIN_ACL_DOMAIN); + + writer.Read("/Root/PQ/account1/topic1", "user1", "topic1@" BUILTIN_ACL_DOMAIN, false, false, false, true); + + writer.Write("Root/PQ/account1/topic1", {"valuevaluevalue2"}, false, "topic1@" BUILTIN_ACL_DOMAIN); //TODO /Root remove + + writer.Read("Root/PQ/account1/topic1", "user1", "topic1@" BUILTIN_ACL_DOMAIN, false, false, false, true); + } +} +Y_UNIT_TEST_SUITE(TPersQueueTest) { + + Y_UNIT_TEST(SetupLockSession2) { + TTestServer server(false); + server.GrpcServerOptions.SetMaxMessageSize(130_MB); + server.StartServer(); + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + server.EnableLogs({ NKikimrServices::PERSQUEUE }, NActors::NLog::PRI_INFO); + server.AnnoyingClient->CreateTopic("rt3.dc1--acc--topic1", 1); + server.AnnoyingClient->CreateTopic("rt3.dc2--acc--topic1", 1); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + server.AnnoyingClient->CreateConsumer("user"); + + auto writer = MakeDataWriter(server); + + TTestPQLib PQLib(server); + auto [consumer, startResult] = PQLib.CreateConsumer({"acc/topic1"}, "user", 1, true); + Cerr << startResult.Response << "\n"; + for (ui32 i = 0; i < 2; ++i) { + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << "Response: " << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetTopic() == "rt3.dc1--acc--topic1" || msg.GetValue().Response.GetLock().GetTopic() == "rt3.dc2--acc--topic1"); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetPartition() == 0); + } + auto msg = consumer->GetNextMessage(); + UNIT_ASSERT(!msg.Wait(TDuration::Seconds(1))); + server.AnnoyingClient->AlterTopic("rt3.dc2--acc--topic1", 2); + msg.Wait(); + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetTopic() == "rt3.dc2--acc--topic1"); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetPartition() == 1); + } + + + + Y_UNIT_TEST(SetupLockSession) { + TTestServer server; + + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + server.EnableLogs({ NKikimrServices::PERSQUEUE }, NActors::NLog::PRI_INFO); + + server.AnnoyingClient->CreateTopic("rt3.dc1--acc--topic1", 1); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + server.AnnoyingClient->CreateConsumer("user"); + + auto writer = MakeDataWriter(server); + + std::shared_ptr<grpc::Channel> Channel_; + std::unique_ptr<NKikimrClient::TGRpcServer::Stub> Stub_; + std::unique_ptr<NPersQueue::PersQueueService::Stub> StubP_; + + Channel_ = grpc::CreateChannel("localhost:" + ToString(server.GrpcPort), grpc::InsecureChannelCredentials()); + Stub_ = NKikimrClient::TGRpcServer::NewStub(Channel_); + + ui64 proxyCookie = 0; + + { + grpc::ClientContext context; + NKikimrClient::TChooseProxyRequest request; + NKikimrClient::TResponse response; + auto status = Stub_->ChooseProxy(&context, request, &response); + UNIT_ASSERT(status.ok()); + Cerr << response << "\n"; + UNIT_ASSERT(response.GetStatus() == NMsgBusProxy::MSTATUS_OK); + proxyCookie = response.GetProxyCookie(); + Channel_ = grpc::CreateChannel("[" + response.GetProxyName() + "]:" + ToString(server.GrpcPort), grpc::InsecureChannelCredentials()); + StubP_ = NPersQueue::PersQueueService::NewStub(Channel_); + } + + + grpc::ClientContext rcontext; + auto readStream = StubP_->ReadSession(&rcontext); + UNIT_ASSERT(readStream); + + // init read session + { + TReadRequest req; + TReadResponse resp; + + req.MutableInit()->AddTopics("acc/topic1"); + + req.MutableInit()->SetClientId("user"); + req.MutableInit()->SetClientsideLocksAllowed(true); + req.MutableInit()->SetProxyCookie(proxyCookie); + req.MutableInit()->SetProtocolVersion(TReadRequest::ReadParamsInInit); + req.MutableInit()->SetMaxReadMessagesCount(3); + + if (!readStream->Write(req)) { + ythrow yexception() << "write fail"; + } + UNIT_ASSERT(readStream->Read(&resp)); + UNIT_ASSERT(resp.HasInit()); + //send some reads + req.Clear(); + req.MutableRead(); + for (ui32 i = 0; i < 10; ++i) { + if (!readStream->Write(req)) { + ythrow yexception() << "write fail"; + } + } + } + + { + TReadRequest req; + TReadResponse resp; + + //lock partition + UNIT_ASSERT(readStream->Read(&resp)); + UNIT_ASSERT(resp.HasLock()); + UNIT_ASSERT_VALUES_EQUAL(resp.GetLock().GetTopic(), "rt3.dc1--acc--topic1"); + UNIT_ASSERT(resp.GetLock().GetPartition() == 0); + + req.Clear(); + req.MutableStartRead()->SetTopic(resp.GetLock().GetTopic()); + req.MutableStartRead()->SetPartition(resp.GetLock().GetPartition()); + req.MutableStartRead()->SetReadOffset(10); + req.MutableStartRead()->SetGeneration(resp.GetLock().GetGeneration()); + if (!readStream->Write(req)) { + ythrow yexception() << "write fail"; + } + + } + + //Write some data + writer.Write("acc/topic1", "valuevaluevalue1"); + Sleep(TDuration::Seconds(15)); //force wait data + writer.Write("acc/topic1", "valuevaluevalue2"); + writer.Write("acc/topic1", "valuevaluevalue3"); + writer.Write("acc/topic1", "valuevaluevalue4"); + + //check read results + TReadResponse resp; + for (ui32 i = 10; i < 16; ++i) { + UNIT_ASSERT(readStream->Read(&resp)); + UNIT_ASSERT_C(resp.HasBatchedData(), resp); + UNIT_ASSERT(resp.GetBatchedData().PartitionDataSize() == 1); + UNIT_ASSERT(resp.GetBatchedData().GetPartitionData(0).BatchSize() == 1); + UNIT_ASSERT(resp.GetBatchedData().GetPartitionData(0).GetBatch(0).MessageDataSize() == 1); + UNIT_ASSERT(resp.GetBatchedData().GetPartitionData(0).GetBatch(0).GetMessageData(0).GetOffset() == i); + } + //TODO: restart here readSession and read from position 10 + { + TReadRequest req; + TReadResponse resp; + + req.MutableCommit()->AddCookie(1); + + if (!readStream->Write(req)) { + ythrow yexception() << "write fail"; + } + UNIT_ASSERT(readStream->Read(&resp)); + UNIT_ASSERT(resp.HasCommit()); + } + } + + + void SetupWriteSessionImpl(bool rr) { + TTestServer server(PQSettings(0, 2, rr)); + + server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY }); + + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 10); + auto writer = MakeDataWriter(server); + + ui32 p = writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue1"); + + server.AnnoyingClient->AlterTopic(DEFAULT_TOPIC_NAME, 15); + + ui32 pp = writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue2"); + UNIT_ASSERT_VALUES_EQUAL(p, pp); + + writer.WriteBatch(SHORT_TOPIC_NAME, {"1", "2", "3", "4", "5"}); + + writer.Write("topic2", "valuevaluevalue1", true); + + p = writer.InitSession("sid1", 2, true); + pp = writer.InitSession("sid1", 0, true); + + UNIT_ASSERT(p = pp); + UNIT_ASSERT(p == 1); + + { + p = writer.InitSession("sidx", 0, true); + pp = writer.InitSession("sidx", 0, true); + + UNIT_ASSERT(p == pp); + } + + writer.InitSession("sid1", 3, false); + + //check round robin; + TMap<ui32, ui32> ss; + for (ui32 i = 0; i < 15*5; ++i) { + ss[writer.InitSession("sid_rand_" + ToString<ui32>(i), 0, true)]++; + } + for (auto &s : ss) { + Cerr << s.first << " " << s.second << "\n"; + if (rr) { + UNIT_ASSERT(s.second >= 4 && s.second <= 6); + } + } + } + + Y_UNIT_TEST(SetupWriteSession) { + SetupWriteSessionImpl(false); + SetupWriteSessionImpl(true); + } + + Y_UNIT_TEST(SetupWriteSessionOnDisabledCluster) { + TTestServer server; + server.EnableLogs({ NKikimrServices::PERSQUEUE, NKikimrServices::PQ_WRITE_PROXY}); + + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 10); + + auto writer = MakeDataWriter(server); + + server.AnnoyingClient->DisableDC(); + + Sleep(TDuration::Seconds(5)); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue1", true); + } + + Y_UNIT_TEST(CloseActiveWriteSessionOnClusterDisable) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 10); + server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY }); + + auto writer = MakeDataWriter(server); + + TTestPQLib PQLib(server); + auto [producer, res] = PQLib.CreateProducer(SHORT_TOPIC_NAME, "123", {}, ECodec::RAW); + + NThreading::TFuture<NPersQueue::TError> isDead = producer->IsDead(); + server.AnnoyingClient->DisableDC(); + isDead.Wait(); + UNIT_ASSERT_EQUAL(isDead.GetValue().GetCode(), NPersQueue::NErrorCode::CLUSTER_DISABLED); + } + + Y_UNIT_TEST(BadSids) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 10); + server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY }); + + auto writer = MakeDataWriter(server); + TTestPQLib PQLib(server); + auto runSidTest = [&](const TString& srcId, bool shouldFail = true) { + auto[producer, res] = PQLib.CreateProducer(SHORT_TOPIC_NAME, srcId); + if (shouldFail) { + UNIT_ASSERT(res.Response.HasError()); + } else { + UNIT_ASSERT(res.Response.HasInit()); + } + }; + + runSidTest("base64:a***"); + runSidTest("base64:aa=="); + runSidTest("base64:a"); + runSidTest("base64:aa", false); + } + + Y_UNIT_TEST(ReadFromSeveralPartitions) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 10); + auto writer = MakeDataWriter(server, "source1"); + + std::shared_ptr<grpc::Channel> Channel_; + std::unique_ptr<NKikimrClient::TGRpcServer::Stub> Stub_; + std::unique_ptr<NPersQueue::PersQueueService::Stub> StubP_; + + Channel_ = grpc::CreateChannel("localhost:" + ToString(server.GrpcPort), grpc::InsecureChannelCredentials()); + Stub_ = NKikimrClient::TGRpcServer::NewStub(Channel_); + + ui64 proxyCookie = 0; + + { + grpc::ClientContext context; + NKikimrClient::TChooseProxyRequest request; + NKikimrClient::TResponse response; + auto status = Stub_->ChooseProxy(&context, request, &response); + UNIT_ASSERT(status.ok()); + Cerr << response << "\n"; + UNIT_ASSERT(response.GetStatus() == NMsgBusProxy::MSTATUS_OK); + proxyCookie = response.GetProxyCookie(); + Channel_ = grpc::CreateChannel("[" + response.GetProxyName() + "]:" + ToString(server.GrpcPort), grpc::InsecureChannelCredentials()); + StubP_ = NPersQueue::PersQueueService::NewStub(Channel_); + } + + + //Write some data + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue1"); + + auto writer2 = MakeDataWriter(server, "source2"); + writer2.Write(SHORT_TOPIC_NAME, "valuevaluevalue2"); + + grpc::ClientContext rcontext; + auto readStream = StubP_->ReadSession(&rcontext); + UNIT_ASSERT(readStream); + + // init read session + { + TReadRequest req; + TReadResponse resp; + + req.MutableInit()->AddTopics(SHORT_TOPIC_NAME); + + req.MutableInit()->SetClientId("user"); + req.MutableInit()->SetProxyCookie(proxyCookie); + req.MutableInit()->SetProtocolVersion(TReadRequest::ReadParamsInInit); + req.MutableInit()->SetMaxReadMessagesCount(1000); + + if (!readStream->Write(req)) { + ythrow yexception() << "write fail"; + } + UNIT_ASSERT(readStream->Read(&resp)); + UNIT_ASSERT(resp.HasInit()); + + //send some reads + Sleep(TDuration::Seconds(5)); + for (ui32 i = 0; i < 10; ++i) { + req.Clear(); + req.MutableRead(); + + if (!readStream->Write(req)) { + ythrow yexception() << "write fail"; + } + } + } + + //check read results + TReadResponse resp; + for (ui32 i = 0; i < 1; ++i) { + UNIT_ASSERT(readStream->Read(&resp)); + UNIT_ASSERT_C(resp.HasBatchedData(), resp); + UNIT_ASSERT(resp.GetBatchedData().PartitionDataSize() == 2); + } + } + + + void SetupReadSessionTest(bool useBatching) { + TTestServer server; + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + server.AnnoyingClient->CreateTopic("rt3.dc2--topic1", 2); + + auto writer = MakeDataWriter(server, "source1"); + + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue0"); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue1"); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue2"); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue3"); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue4"); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue5"); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue6"); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue7"); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue8"); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue9"); + + writer.Read(SHORT_TOPIC_NAME, "user", "", false, false, useBatching); + } + + Y_UNIT_TEST(SetupReadSession) { + SetupReadSessionTest(false); + } + + Y_UNIT_TEST(SetupReadSessionWithBatching) { + SetupReadSessionTest(true); + } + + void ClosesSessionOnReadSettingsChangeTest(bool initReadSettingsInInitRequest) { + TTestServer server; + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + server.AnnoyingClient->CreateTopic("rt3.dc2--topic1", 2); + + auto writer = MakeDataWriter(server, "source1"); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue0"); + + // Reading code + std::shared_ptr<grpc::Channel> Channel_; + std::unique_ptr<NKikimrClient::TGRpcServer::Stub> Stub_; + std::unique_ptr<NPersQueue::PersQueueService::Stub> StubP_; + + Channel_ = grpc::CreateChannel("localhost:" + ToString(server.GrpcPort), grpc::InsecureChannelCredentials()); + Stub_ = NKikimrClient::TGRpcServer::NewStub(Channel_); + + ui64 proxyCookie = 0; + + { + grpc::ClientContext context; + NKikimrClient::TChooseProxyRequest request; + NKikimrClient::TResponse response; + auto status = Stub_->ChooseProxy(&context, request, &response); + UNIT_ASSERT(status.ok()); + Cerr << response << "\n"; + UNIT_ASSERT(response.GetStatus() == NMsgBusProxy::MSTATUS_OK); + proxyCookie = response.GetProxyCookie(); + Channel_ = grpc::CreateChannel("[" + response.GetProxyName() + "]:" + ToString(server.GrpcPort), grpc::InsecureChannelCredentials()); + StubP_ = NPersQueue::PersQueueService::NewStub(Channel_); + } + + grpc::ClientContext rcontext; + auto readStream = StubP_->ReadSession(&rcontext); + UNIT_ASSERT(readStream); + + // init read session + { + TReadRequest req; + TReadResponse resp; + + req.MutableInit()->AddTopics(SHORT_TOPIC_NAME); + + req.MutableInit()->SetClientId("user"); + req.MutableInit()->SetProxyCookie(proxyCookie); + if (initReadSettingsInInitRequest) { + req.MutableInit()->SetProtocolVersion(TReadRequest::ReadParamsInInit); + req.MutableInit()->SetMaxReadMessagesCount(1); + } + + if (!readStream->Write(req)) { + ythrow yexception() << "write fail"; + } + UNIT_ASSERT(readStream->Read(&resp)); + UNIT_ASSERT(resp.HasInit()); + + if (!initReadSettingsInInitRequest) { + // send first read + req.Clear(); + req.MutableRead()->SetMaxCount(1); + if (!readStream->Write(req)) { + ythrow yexception() << "write fail"; + } + UNIT_ASSERT(readStream->Read(&resp)); + UNIT_ASSERT_C(resp.HasData(), resp); + } + + // change settings + req.Clear(); + req.MutableRead()->SetMaxCount(42); + if (!readStream->Write(req)) { + ythrow yexception() << "write fail"; + } + UNIT_ASSERT(readStream->Read(&resp)); + UNIT_ASSERT(resp.HasError()); + } + } + + Y_UNIT_TEST(WriteSessionClose) { + + TTestServer server; + server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY }); + + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + server.AnnoyingClient->CreateTopic("rt3.dc2--topic1", 2); + + auto writer = MakeDataWriter(server, "source1"); + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue0"); + + // Reading code + std::shared_ptr<grpc::Channel> Channel_; + std::unique_ptr<NKikimrClient::TGRpcServer::Stub> Stub_; + std::unique_ptr<NPersQueue::PersQueueService::Stub> StubP_; + + Channel_ = grpc::CreateChannel("localhost:" + ToString(server.GrpcPort), grpc::InsecureChannelCredentials()); + Stub_ = NKikimrClient::TGRpcServer::NewStub(Channel_); + + ui64 proxyCookie = 0; + + { + grpc::ClientContext context; + NKikimrClient::TChooseProxyRequest request; + NKikimrClient::TResponse response; + auto status = Stub_->ChooseProxy(&context, request, &response); + UNIT_ASSERT(status.ok()); + Cerr << response << "\n"; + UNIT_ASSERT(response.GetStatus() == NMsgBusProxy::MSTATUS_OK); + proxyCookie = response.GetProxyCookie(); + Channel_ = grpc::CreateChannel("[" + response.GetProxyName() + "]:" + ToString(server.GrpcPort), grpc::InsecureChannelCredentials()); + StubP_ = NPersQueue::PersQueueService::NewStub(Channel_); + } + + // init write session + for (ui32 i = 0; i < 2; ++i){ + grpc::ClientContext rcontext; + + auto writeStream = StubP_->WriteSession(&rcontext); + UNIT_ASSERT(writeStream); + + TWriteRequest req; + + req.MutableInit()->SetTopic(SHORT_TOPIC_NAME); + + req.MutableInit()->SetSourceId("user"); + req.MutableInit()->SetProxyCookie(proxyCookie); + if (i == 0) + continue; + if (!writeStream->Write(req)) { + ythrow yexception() << "write fail"; + } + } + } + + Y_UNIT_TEST(ClosesSessionOnReadSettingsChange) { + ClosesSessionOnReadSettingsChangeTest(false); + } + + Y_UNIT_TEST(ClosesSessionOnReadSettingsChangeWithInit) { + ClosesSessionOnReadSettingsChangeTest(true); + } + + Y_UNIT_TEST(WriteExisting) { + TTestServer server; + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD }); + + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + + { + THolder<NMsgBusProxy::TBusPersQueue> request = TRequestDescribePQ().GetRequest({}); + + NKikimrClient::TResponse response; + + auto channel = grpc::CreateChannel("localhost:"+ToString(server.GrpcPort), grpc::InsecureChannelCredentials()); + auto stub(NKikimrClient::TGRpcServer::NewStub(channel)); + grpc::ClientContext context; + auto status = stub->PersQueueRequest(&context, request->Record, &response); + + UNIT_ASSERT(status.ok()); + } + + server.AnnoyingClient->WriteToPQ( + DEFAULT_TOPIC_NAME, 1, "abacaba", 1, "valuevaluevalue1", "", ETransport::GRpc + ); + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 1, "abacaba", 2, "valuevaluevalue1", "", ETransport::GRpc); + } + + Y_UNIT_TEST(WriteExistingBigValue) { + TTestServer server(PQSettings(0, 2).SetDomainName("Root")); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2, 8_MB, 86400, 100000); + + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD }); + + TInstant now(Now()); + + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 1, "abacaba", 1, TString(1000000, 'a')); + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 1, "abacaba", 2, TString(1, 'a')); + UNIT_ASSERT(TInstant::Now() - now > TDuration::MilliSeconds(5990)); //speed limit is 200kb/s and burst is 200kb, so to write 1mb it will take at least 4 seconds + } + + Y_UNIT_TEST(WriteEmptyData) { + TTestServer server(PQSettings(0, 2).SetDomainName("Root")); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + + // empty data and sourecId + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 1, "", 1, "", "", ETransport::MsgBus, NMsgBusProxy::MSTATUS_ERROR); + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 1, "a", 1, "", "", ETransport::MsgBus, NMsgBusProxy::MSTATUS_ERROR); + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 1, "", 1, "a", "", ETransport::MsgBus, NMsgBusProxy::MSTATUS_ERROR); + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 1, "a", 1, "a", "", ETransport::MsgBus, NMsgBusProxy::MSTATUS_OK); + } + + + Y_UNIT_TEST(WriteNonExistingPartition) { + TTestServer server(PQSettings(0, 2).SetDomainName("Root")); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD }); + + server.AnnoyingClient->WriteToPQ( + DEFAULT_TOPIC_NAME, 100500, "abacaba", 1, "valuevaluevalue1", "", + ETransport::MsgBus, NMsgBusProxy::MSTATUS_ERROR, NMsgBusProxy::MSTATUS_ERROR + ); + } + + Y_UNIT_TEST(WriteNonExistingTopic) { + TTestServer server(PQSettings(0, 2).SetDomainName("Root")); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD }); + + server.AnnoyingClient->WriteToPQ("rt3.dc1--topic1000", 1, "abacaba", 1, "valuevaluevalue1", "", ETransport::MsgBus, NMsgBusProxy::MSTATUS_ERROR, NMsgBusProxy::MSTATUS_ERROR); + } + + Y_UNIT_TEST(SchemeshardRestart) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + TString topic2 = "rt3.dc1--topic2"; + server.AnnoyingClient->CreateTopic(topic2, 2); + + // force topic1 into cache and establish pipe from cache to schemeshard + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 1, "abacaba", 1, "valuevaluevalue1"); + + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD, NKikimrServices::PQ_METACACHE }); + + server.AnnoyingClient->RestartSchemeshard(server.CleverServer->GetRuntime()); + server.AnnoyingClient->WriteToPQ(topic2, 1, "abacaba", 1, "valuevaluevalue1"); + } + + Y_UNIT_TEST(WriteAfterAlter) { + TTestServer server(PQSettings(0).SetDomainName("Root").SetNodeCount(2)); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD, NKikimrServices::PQ_METACACHE }); + + + server.AnnoyingClient->WriteToPQ( + DEFAULT_TOPIC_NAME, 5, "abacaba", 1, "valuevaluevalue1", "", + ETransport::MsgBus, NMsgBusProxy::MSTATUS_ERROR, NMsgBusProxy::MSTATUS_ERROR + ); + + server.AnnoyingClient->AlterTopic(DEFAULT_TOPIC_NAME, 10); + Sleep(TDuration::Seconds(1)); + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 5, "abacaba", 1, "valuevaluevalue1"); + server.AnnoyingClient->WriteToPQ( + DEFAULT_TOPIC_NAME, 15, "abacaba", 1, "valuevaluevalue1", "", + ETransport::MsgBus, NMsgBusProxy::MSTATUS_ERROR, NMsgBusProxy::MSTATUS_ERROR + ); + + server.AnnoyingClient->AlterTopic(DEFAULT_TOPIC_NAME, 20); + Sleep(TDuration::Seconds(1)); + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 5, "abacaba", 1, "valuevaluevalue1"); + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 15, "abacaba", 1, "valuevaluevalue1"); + } + + Y_UNIT_TEST(Delete) { + TTestServer server(PQSettings(0).SetDomainName("Root").SetNodeCount(2)); + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD }); + + // Delete non-existing + server.AnnoyingClient->DeleteTopic2(DEFAULT_TOPIC_NAME, NPersQueue::NErrorCode::UNKNOWN_TOPIC); + + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + + // Delete existing + server.AnnoyingClient->DeleteTopic2(DEFAULT_TOPIC_NAME); + + // Double delete - "What Is Dead May Never Die" + server.AnnoyingClient->DeleteTopic2(DEFAULT_TOPIC_NAME, NPersQueue::NErrorCode::UNKNOWN_TOPIC); + + // Resurrect deleted topic + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + server.AnnoyingClient->DeleteTopic2(DEFAULT_TOPIC_NAME); + } + + Y_UNIT_TEST(WriteAfterDelete) { + TTestServer server(PQSettings(0).SetDomainName("Root").SetNodeCount(2)); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 3); + + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 1, "abacaba", 1, "valuevaluevalue1"); + + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD, NKikimrServices::PQ_METACACHE }); + + server.AnnoyingClient->DeleteTopic2(DEFAULT_TOPIC_NAME); + server.AnnoyingClient->WriteToPQ( + DEFAULT_TOPIC_NAME, 1, "abacaba", 2, "valuevaluevalue1", "", + ETransport::MsgBus, NMsgBusProxy::MSTATUS_ERROR, NMsgBusProxy::MSTATUS_ERROR + ); + server.AnnoyingClient->WriteToPQ( + DEFAULT_TOPIC_NAME, 2, "abacaba", 1, "valuevaluevalue1", "", + ETransport::MsgBus, NMsgBusProxy::MSTATUS_ERROR, NMsgBusProxy::MSTATUS_ERROR + ); + } + + Y_UNIT_TEST(WriteAfterCreateDeleteCreate) { + TTestServer server(PQSettings(0).SetDomainName("Root").SetNodeCount(2)); + TString topic = "rt3.dc1--johnsnow"; + server.AnnoyingClient->CreateTopic(topic, 2); + + server.AnnoyingClient->WriteToPQ(topic, 1, "abacaba", 1, "valuevaluevalue1"); + server.AnnoyingClient->WriteToPQ(topic, 3, "abacaba", 1, "valuevaluevalue1", "", ETransport::MsgBus, NMsgBusProxy::MSTATUS_ERROR, NMsgBusProxy::MSTATUS_ERROR); + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD, NKikimrServices::PQ_METACACHE }); + + server.AnnoyingClient->DeleteTopic2(topic); + + server.AnnoyingClient->CreateTopic(topic, 4); + + // Write to topic, cache must be updated by CreateTopic + server.AnnoyingClient->WriteToPQ(topic, 1, "abacaba", 1, "valuevaluevalue1"); + // Write to partition that didn't exist in the old topic + server.AnnoyingClient->WriteToPQ(topic, 3, "abacaba", 1, "valuevaluevalue1"); + } + + Y_UNIT_TEST(GetOffsetsAfterDelete) { + TTestServer server(PQSettings(0).SetDomainName("Root").SetNodeCount(2)); + TString topic2 = "rt3.dc1--topic2"; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 3); + server.AnnoyingClient->CreateTopic(topic2, 3); + + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 1, "abacaba", 1, "valuevaluevalue1"); + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD, NKikimrServices::PQ_METACACHE }); + + server.AnnoyingClient->DeleteTopic2(DEFAULT_TOPIC_NAME); + + // Get offsets from deleted topic + server.AnnoyingClient->GetPartOffset( { + {DEFAULT_TOPIC_NAME, {1,2}} + }, 0, 0, false); + + // Get offsets from multiple topics + server.AnnoyingClient->GetPartOffset( { + {DEFAULT_TOPIC_NAME, {1,2}}, + {topic2, {1,2}}, + }, 0, 0, false); + } + + + Y_UNIT_TEST(GetOffsetsAfterCreateDeleteCreate) { + TTestServer server(PQSettings(0).SetDomainName("Root").SetNodeCount(2)); + TString topic2 = "rt3.dc1--topic2"; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 3); + server.AnnoyingClient->CreateTopic(topic2, 3); + + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 1, "abacaba", 1, "valuevaluevalue1"); + + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD, NKikimrServices::PQ_METACACHE }); + + server.AnnoyingClient->DeleteTopic2(DEFAULT_TOPIC_NAME); + Sleep(TDuration::Seconds(1)); + + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 4); + Sleep(TDuration::Seconds(1)); + + // Get offsets from multiple topics + server.AnnoyingClient->GetPartOffset( { + {DEFAULT_TOPIC_NAME, {1,2}}, + {topic2, {1}}, + }, 3, 0, true); + } + + Y_UNIT_TEST(BigRead) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1, 8_MB, 86400, 20000000, "user", 2000000); + + server.EnableLogs( { NKikimrServices::FLAT_TX_SCHEMESHARD }); + + TString value(1_MB, 'x'); + for (ui32 i = 0; i < 32; ++i) + server.AnnoyingClient->WriteToPQ({DEFAULT_TOPIC_NAME, 0, "source1", i}, value); + + // trying to read small PQ messages in a big messagebus event + auto info = server.AnnoyingClient->ReadFromPQ({DEFAULT_TOPIC_NAME, 0, 0, 32, "user"}, 23, "", NMsgBusProxy::MSTATUS_OK); //will read 21mb + UNIT_ASSERT_VALUES_EQUAL(info.BlobsFromDisk, 0); + UNIT_ASSERT_VALUES_EQUAL(info.BlobsFromCache, 4); + + TInstant now(TInstant::Now()); + info = server.AnnoyingClient->ReadFromPQ({DEFAULT_TOPIC_NAME, 0, 0, 32, "user"}, 23, "", NMsgBusProxy::MSTATUS_OK); //will read 21mb + TDuration dur = TInstant::Now() - now; + UNIT_ASSERT_C(dur > TDuration::Seconds(7) && dur < TDuration::Seconds(20), "dur = " << dur); //speed limit is 2000kb/s and burst is 2000kb, so to read 24mb it will take at least 11 seconds + + server.AnnoyingClient->GetPartStatus({}, 1, true); + + } + + // expects that L2 size is 32Mb + Y_UNIT_TEST(Cache) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1, 8_MB); + + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD }); + + TString value(1_MB, 'x'); + for (ui32 i = 0; i < 32; ++i) + server.AnnoyingClient->WriteToPQ({DEFAULT_TOPIC_NAME, 0, "source1", i}, value); + + auto info0 = server.AnnoyingClient->ReadFromPQ({DEFAULT_TOPIC_NAME, 0, 0, 16, "user"}, 16); + auto info16 = server.AnnoyingClient->ReadFromPQ({DEFAULT_TOPIC_NAME, 0, 16, 16, "user"}, 16); + + UNIT_ASSERT_VALUES_EQUAL(info0.BlobsFromCache, 3); + UNIT_ASSERT_VALUES_EQUAL(info16.BlobsFromCache, 2); + UNIT_ASSERT_VALUES_EQUAL(info0.BlobsFromDisk + info16.BlobsFromDisk, 0); + + for (ui32 i = 0; i < 8; ++i) + server.AnnoyingClient->WriteToPQ({DEFAULT_TOPIC_NAME, 0, "source1", 32+i}, value); + + info0 = server.AnnoyingClient->ReadFromPQ({DEFAULT_TOPIC_NAME, 0, 0, 16, "user"}, 16); + info16 = server.AnnoyingClient->ReadFromPQ({DEFAULT_TOPIC_NAME, 0, 16, 16, "user"}, 16); + + ui32 fromDisk = info0.BlobsFromDisk + info16.BlobsFromDisk; + ui32 fromCache = info0.BlobsFromCache + info16.BlobsFromCache; + UNIT_ASSERT(fromDisk > 0); + UNIT_ASSERT(fromDisk < 5); + UNIT_ASSERT(fromCache > 0); + UNIT_ASSERT(fromCache < 5); + } + + Y_UNIT_TEST(CacheHead) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1, 6_MB); + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD }); + + ui64 seqNo = 0; + for (ui32 blobSizeKB = 256; blobSizeKB < 4096; blobSizeKB *= 2) { + static const ui32 maxEventKB = 24_KB; + ui32 blobSize = blobSizeKB * 1_KB; + ui32 count = maxEventKB / blobSizeKB; + count -= count%2; + ui32 half = count/2; + + ui64 offset = seqNo; + TString value(blobSize, 'a'); + for (ui32 i = 0; i < count; ++i) + server.AnnoyingClient->WriteToPQ({DEFAULT_TOPIC_NAME, 0, "source1", seqNo++}, value); + + auto info_half1 = server.AnnoyingClient->ReadFromPQ({DEFAULT_TOPIC_NAME, 0, offset, half, "user1"}, half); + auto info_half2 = server.AnnoyingClient->ReadFromPQ({DEFAULT_TOPIC_NAME, 0, offset, half, "user1"}, half); + + UNIT_ASSERT(info_half1.BlobsFromCache > 0); + UNIT_ASSERT(info_half2.BlobsFromCache > 0); + UNIT_ASSERT_VALUES_EQUAL(info_half1.BlobsFromDisk, 0); + UNIT_ASSERT_VALUES_EQUAL(info_half2.BlobsFromDisk, 0); + } + } + + Y_UNIT_TEST(SameOffset) { + TTestServer server; + TString topic2 = "rt3.dc1--topic2"; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1, 6_MB); + server.AnnoyingClient->CreateTopic(topic2, 1, 6_MB); + + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD }); + + ui32 valueSize = 128; + TString value1(valueSize, 'a'); + TString value2(valueSize, 'b'); + server.AnnoyingClient->WriteToPQ({DEFAULT_TOPIC_NAME, 0, "source1", 0}, value1); + server.AnnoyingClient->WriteToPQ({topic2, 0, "source1", 0}, value2); + + // avoid reading from head + TString mb(1_MB, 'x'); + for (ui32 i = 1; i < 16; ++i) { + server.AnnoyingClient->WriteToPQ({DEFAULT_TOPIC_NAME, 0, "source1", i}, mb); + server.AnnoyingClient->WriteToPQ({topic2, 0, "source1", i}, mb); + } + + auto info1 = server.AnnoyingClient->ReadFromPQ({DEFAULT_TOPIC_NAME, 0, 0, 1, "user1"}, 1); + auto info2 = server.AnnoyingClient->ReadFromPQ({topic2, 0, 0, 1, "user1"}, 1); + + UNIT_ASSERT_VALUES_EQUAL(info1.BlobsFromCache, 1); + UNIT_ASSERT_VALUES_EQUAL(info2.BlobsFromCache, 1); + UNIT_ASSERT_VALUES_EQUAL(info1.Values.size(), 1); + UNIT_ASSERT_VALUES_EQUAL(info2.Values.size(), 1); + UNIT_ASSERT_VALUES_EQUAL(info1.Values[0].size(), valueSize); + UNIT_ASSERT_VALUES_EQUAL(info2.Values[0].size(), valueSize); + UNIT_ASSERT(info1.Values[0] == value1); + UNIT_ASSERT(info2.Values[0] == value2); + } + + + Y_UNIT_TEST(FetchRequest) { + TTestServer server; + TString topic2 = "rt3.dc1--topic2"; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 10); + server.AnnoyingClient->CreateTopic(topic2, 10); + + ui32 valueSize = 128; + TString value1(valueSize, 'a'); + TString value2(valueSize, 'b'); + server.AnnoyingClient->WriteToPQ({topic2, 5, "source1", 0}, value2); + server.AnnoyingClient->WriteToPQ({DEFAULT_TOPIC_NAME, 1, "source1", 0}, value1); + server.AnnoyingClient->WriteToPQ({DEFAULT_TOPIC_NAME, 1, "source1", 1}, value2); + + server.EnableLogs({ NKikimrServices::FLAT_TX_SCHEMESHARD }); + TInstant tm(TInstant::Now()); + server.AnnoyingClient->FetchRequestPQ( + { {topic2, 5, 0, 400}, {DEFAULT_TOPIC_NAME, 1, 0, 400}, {DEFAULT_TOPIC_NAME, 3, 0, 400} }, + 400, 1000000 + ); + UNIT_ASSERT((TInstant::Now() - tm).Seconds() < 1); + tm = TInstant::Now(); + server.AnnoyingClient->FetchRequestPQ({{topic2, 5, 1, 400}}, 400, 5000); + UNIT_ASSERT((TInstant::Now() - tm).Seconds() > 2); + server.AnnoyingClient->FetchRequestPQ( + { {topic2, 5, 0, 400}, {DEFAULT_TOPIC_NAME, 1, 0, 400}, {DEFAULT_TOPIC_NAME, 3, 0, 400} }, + 1, 1000000 + ); + server.AnnoyingClient->FetchRequestPQ( + { {topic2, 5, 500, 400}, {topic2, 4, 0, 400}, {DEFAULT_TOPIC_NAME, 1, 0, 400} }, + 400, 1000000 + ); + } + + Y_UNIT_TEST(ChooseProxy) { + TTestServer server; + server.AnnoyingClient->ChooseProxy(ETransport::GRpc); + } + + + Y_UNIT_TEST(Init) { + TTestServer server(PQSettings(0).SetDomainName("Root").SetNodeCount(2)); + TString topic2 = "rt3.dc1--topic2"; + TString topic3 = "rt3.dc1--topic3"; + + if (!true) { + server.EnableLogs({ + NKikimrServices::FLAT_TX_SCHEMESHARD, + NKikimrServices::TX_DATASHARD, + NKikimrServices::HIVE, + NKikimrServices::PERSQUEUE, + NKikimrServices::TABLET_MAIN, + NKikimrServices::BS_PROXY_DISCOVER, + NKikimrServices::PIPE_CLIENT, + NKikimrServices::PQ_METACACHE }); + } + + server.AnnoyingClient->DescribeTopic({}); + server.AnnoyingClient->TestCase({}, 0, 0, true); + + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 10); + server.AnnoyingClient->AlterTopic(DEFAULT_TOPIC_NAME, 20); + server.AnnoyingClient->CreateTopic(topic2, 25); + + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 5, "abacaba", 1, "valuevaluevalue1"); + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 5, "abacaba", 2, "valuevaluevalue2"); + server.AnnoyingClient->WriteToPQ(DEFAULT_TOPIC_NAME, 5, "abacabae", 1, "valuevaluevalue3"); + server.AnnoyingClient->ReadFromPQ(DEFAULT_TOPIC_NAME, 5, 0, 10, 3); + + server.AnnoyingClient->SetClientOffsetPQ(DEFAULT_TOPIC_NAME, 5, 2); + + server.AnnoyingClient->TestCase({{DEFAULT_TOPIC_NAME, {5}}}, 1, 1, true); + server.AnnoyingClient->TestCase({{DEFAULT_TOPIC_NAME, {0}}}, 1, 0, true); + server.AnnoyingClient->TestCase({{DEFAULT_TOPIC_NAME, {}}}, 20, 1, true); + server.AnnoyingClient->TestCase({{DEFAULT_TOPIC_NAME, {5, 5}}}, 0, 0, false); + server.AnnoyingClient->TestCase({{DEFAULT_TOPIC_NAME, {111}}}, 0, 0, false); + server.AnnoyingClient->TestCase({}, 45, 1, true); + server.AnnoyingClient->TestCase({{topic3, {}}}, 0, 0, false); + server.AnnoyingClient->TestCase({{DEFAULT_TOPIC_NAME, {}}, {topic3, {}}}, 0, 0, false); + server.AnnoyingClient->TestCase({{DEFAULT_TOPIC_NAME, {}}, {topic2, {}}}, 45, 1, true); + server.AnnoyingClient->TestCase({{DEFAULT_TOPIC_NAME, {0, 3, 5}}, {topic2, {1, 4, 6, 8}}}, 7, 1, true); + + server.AnnoyingClient->DescribeTopic({DEFAULT_TOPIC_NAME}); + server.AnnoyingClient->DescribeTopic({topic2}); + server.AnnoyingClient->DescribeTopic({topic2, DEFAULT_TOPIC_NAME}); + server.AnnoyingClient->DescribeTopic({}); + server.AnnoyingClient->DescribeTopic({topic3}, true); + } + + + Y_UNIT_TEST(DescribeAfterDelete) { + TTestServer server(PQSettings(0).SetDomainName("Root").SetNodeCount(2)); + server.EnableLogs({ NKikimrServices::PQ_METACACHE }); + TString topic2 = "rt3.dc1--topic2"; + + server.AnnoyingClient->DescribeTopic({}); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 10); + server.AnnoyingClient->CreateTopic(topic2, 10); + server.AnnoyingClient->DescribeTopic({}); + + server.AnnoyingClient->DeleteTopic2(DEFAULT_TOPIC_NAME); + + server.AnnoyingClient->DescribeTopic({}); + server.AnnoyingClient->GetClientInfo({}, "user", true); + server.AnnoyingClient->GetClientInfo({topic2}, "user", true); + Sleep(TDuration::Seconds(2)); + server.AnnoyingClient->TestCase({{DEFAULT_TOPIC_NAME, {}}, {topic2, {}}}, 10, 0, false); + server.AnnoyingClient->TestCase({{DEFAULT_TOPIC_NAME, {}}, {topic2, {}}, {"rt3.dc1--topic3", {}}}, 10, 0, false); + } + + Y_UNIT_TEST(DescribeAfterDelete2) { + TTestServer server(PQSettings(0).SetDomainName("Root").SetNodeCount(2)); + server.EnableLogs({ NKikimrServices::PQ_METACACHE }); + TString topic3 = "rt3.dc1--topic3"; + + server.AnnoyingClient->CreateTopic(topic3, 10); + Sleep(TDuration::Seconds(3)); //for invalidation of cache + server.AnnoyingClient->TestCase({{"rt3.dc1--topic4", {}}}, 10, 0, false); //will force caching of topic3 + server.AnnoyingClient->DeleteTopic2(topic3); + server.AnnoyingClient->DescribeTopic({topic3}, true); //describe will fail + server.AnnoyingClient->DescribeTopic({topic3}, true); //describe will fail + } + + + void WaitResolveSuccess(TTestServer& server, TString topic, ui32 numParts) { + const TInstant start = TInstant::Now(); + while (true) { + TAutoPtr<NMsgBusProxy::TBusPersQueue> request(new NMsgBusProxy::TBusPersQueue); + auto req = request->Record.MutableMetaRequest(); + auto partOff = req->MutableCmdGetPartitionLocations(); + auto treq = partOff->AddTopicRequest(); + treq->SetTopic(topic); + for (ui32 i = 0; i < numParts; ++i) + treq->AddPartition(i); + + TAutoPtr<NBus::TBusMessage> reply; + NBus::EMessageStatus status = server.AnnoyingClient->SyncCall(request, reply); + UNIT_ASSERT_VALUES_EQUAL(status, NBus::MESSAGE_OK); + const NMsgBusProxy::TBusResponse* response = dynamic_cast<NMsgBusProxy::TBusResponse*>(reply.Get()); + UNIT_ASSERT(response); + if (response->Record.GetStatus() == NMsgBusProxy::MSTATUS_OK) + break; + UNIT_ASSERT(TInstant::Now() - start < ::DEFAULT_DISPATCH_TIMEOUT); + Sleep(TDuration::MilliSeconds(10)); + } + } + + Y_UNIT_TEST(WhenDisableNodeAndCreateTopic_ThenAllPartitionsAreOnOtherNode) { + // Arrange + TTestServer server(PQSettings(0).SetDomainName("Root").SetNodeCount(2)); + server.EnableLogs({ NKikimrServices::HIVE }); + TString unused = "rt3.dc1--unusedtopic"; + // Just to make sure that HIVE has started + server.AnnoyingClient->CreateTopic(unused, 1); + WaitResolveSuccess(server, unused, 1); + + // Act + // Disable node #0 + server.AnnoyingClient->MarkNodeInHive(server.CleverServer->GetRuntime(), 0, false); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 3); + WaitResolveSuccess(server, DEFAULT_TOPIC_NAME, 3); + + // Assert that all partitions are on node #1 + const ui32 node1Id = server.CleverServer->GetRuntime()->GetNodeId(1); + UNIT_ASSERT_VALUES_EQUAL( + server.AnnoyingClient->GetPartLocation({{DEFAULT_TOPIC_NAME, {0, 1}}}, 2, true), + TVector<ui32>({node1Id, node1Id}) + ); + } + + void PrepareForGrpc(TTestServer& server) { + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 2); + server.AnnoyingClient->InitUserRegistry(); + } + void PrepareForFstClass(TFlatMsgBusPQClient& annoyingClient) { + annoyingClient.SetNoConfigMode(); + annoyingClient.FullInit({}); + annoyingClient.InitUserRegistry(); + annoyingClient.MkDir("/Root", "account1"); + annoyingClient.CreateTopicNoLegacy(FC_TOPIC_PATH, 5); + } + + Y_UNIT_TEST(CheckACLForGrpcWrite) { + TTestServer server; + PrepareForGrpc(server); + + auto writer = MakeDataWriter(server, "source1"); + + writer.Write(SHORT_TOPIC_NAME, "valuevaluevalue1", true, "topic1@" BUILTIN_ACL_DOMAIN); + + NACLib::TDiffACL acl; + acl.AddAccess(NACLib::EAccessType::Allow, NACLib::UpdateRow, "topic1@" BUILTIN_ACL_DOMAIN); + server.AnnoyingClient->ModifyACL("/Root/PQ", DEFAULT_TOPIC_NAME, acl.SerializeAsString()); + + Sleep(TDuration::Seconds(2)); + + auto writer2 = MakeDataWriter(server, "source1"); + writer2.Write(SHORT_TOPIC_NAME, "valuevaluevalue1", false, "topic1@" BUILTIN_ACL_DOMAIN); + } + + void PrepareForGrpcNoDC(TFlatMsgBusPQClient& annoyingClient) { + annoyingClient.SetNoConfigMode(); + annoyingClient.FullInit(); + annoyingClient.InitUserRegistry(); + annoyingClient.MkDir("/Root", "account1"); + annoyingClient.MkDir("/Root/PQ", "account1"); + annoyingClient.CreateTopicNoLegacy("/Root/PQ/rt3.db--topic1", 5, false); + annoyingClient.CreateTopicNoLegacy("/Root/PQ/account1/topic1", 5, false, true, {}, {"user1", "user2"}); + annoyingClient.CreateTopicNoLegacy("/Root/account2/topic2", 5); + } + + Y_UNIT_TEST(TestGrpcWriteNoDC) { + TTestServer server(false); + server.ServerSettings.PQConfig.SetTopicsAreFirstClassCitizen(true); + server.ServerSettings.PQConfig.SetRoot("/Rt2/PQ"); + server.ServerSettings.PQConfig.SetDatabase("/Root"); + server.StartServer(); + + UNIT_ASSERT_VALUES_EQUAL(NMsgBusProxy::MSTATUS_OK, + server.AnnoyingClient->AlterUserAttributes("/", "Root", {{"folder_id", "somefolder"}, {"cloud_id", "somecloud"}, {"database_id", "root"}})); + + + PrepareForGrpcNoDC(*server.AnnoyingClient); + auto writer = MakeDataWriter(server, "source1"); + + writer.Write("/Root/account2/topic2", {"valuevaluevalue1"}, true, "topic1@" BUILTIN_ACL_DOMAIN); + writer.Write("/Root/PQ/account1/topic1", {"valuevaluevalue1"}, true, "topic1@" BUILTIN_ACL_DOMAIN); + + NACLib::TDiffACL acl; + acl.AddAccess(NACLib::EAccessType::Allow, NACLib::UpdateRow, "topic1@" BUILTIN_ACL_DOMAIN); + acl.AddAccess(NACLib::EAccessType::Allow, NACLib::SelectRow, "topic1@" BUILTIN_ACL_DOMAIN); + + server.AnnoyingClient->ModifyACL("/Root/account2", "topic2", acl.SerializeAsString()); + server.AnnoyingClient->ModifyACL("/Root/PQ/account1", "topic1", acl.SerializeAsString()); + + Sleep(TDuration::Seconds(5)); + + writer.Write("/Root/account2/topic2", {"valuevaluevalue1"}, false, "topic1@" BUILTIN_ACL_DOMAIN); + + writer.Write("/Root/PQ/account1/topic1", {"valuevaluevalue1"}, false, "topic1@" BUILTIN_ACL_DOMAIN); + writer.Write("/Root/PQ/account1/topic1", {"valuevaluevalue2"}, false, "topic1@" BUILTIN_ACL_DOMAIN); + + writer.Read("/Root/PQ/account1/topic1", "user1", "topic1@" BUILTIN_ACL_DOMAIN, false, false, false, true); + + writer.Write("Root/PQ/account1/topic1", {"valuevaluevalue2"}, false, "topic1@" BUILTIN_ACL_DOMAIN); //TODO /Root remove + + writer.Read("Root/PQ/account1/topic1", "user1", "topic1@" BUILTIN_ACL_DOMAIN, false, false, false, true); + } + + Y_UNIT_TEST(CheckACLForGrpcRead) { + TTestServer server(PQSettings(0, 1)); + PrepareForGrpc(server); + TString topic2 = "rt3.dc1--topic2"; + TString topic2ShortName = "topic2"; + server.AnnoyingClient->CreateTopic( + topic2, 1, 8_MB, 86400, 20000000, "", 200000000, {"user1", "user2"} + ); + server.EnableLogs({NKikimrServices::PERSQUEUE}, NActors::NLog::PRI_INFO); + + server.AnnoyingClient->CreateConsumer("user1"); + server.AnnoyingClient->CreateConsumer("user2"); + server.AnnoyingClient->CreateConsumer("user5"); + server.AnnoyingClient->GrantConsumerAccess("user1", "user2@" BUILTIN_ACL_DOMAIN); + server.AnnoyingClient->GrantConsumerAccess("user1", "user3@" BUILTIN_ACL_DOMAIN); + + server.AnnoyingClient->GrantConsumerAccess("user1", "1@" BUILTIN_ACL_DOMAIN); + server.AnnoyingClient->GrantConsumerAccess("user2", "2@" BUILTIN_ACL_DOMAIN); + server.AnnoyingClient->GrantConsumerAccess("user5", "1@" BUILTIN_ACL_DOMAIN); + server.AnnoyingClient->GrantConsumerAccess("user5", "2@" BUILTIN_ACL_DOMAIN); + + auto writer = MakeDataWriter(server, "source1"); + + NACLib::TDiffACL acl; + acl.AddAccess(NACLib::EAccessType::Allow, NACLib::SelectRow, "1@" BUILTIN_ACL_DOMAIN); + acl.AddAccess(NACLib::EAccessType::Allow, NACLib::SelectRow, "2@" BUILTIN_ACL_DOMAIN); + acl.AddAccess(NACLib::EAccessType::Allow, NACLib::SelectRow, "user1@" BUILTIN_ACL_DOMAIN); + acl.AddAccess(NACLib::EAccessType::Allow, NACLib::SelectRow, "user2@" BUILTIN_ACL_DOMAIN); + server.AnnoyingClient->ModifyACL("/Root/PQ", topic2, acl.SerializeAsString()); + + Sleep(TDuration::Seconds(2)); + + auto ticket1 = "1@" BUILTIN_ACL_DOMAIN; + auto ticket2 = "2@" BUILTIN_ACL_DOMAIN; + + writer.Read(topic2ShortName, "user1", ticket1, false, false, false, true); + + writer.Read(topic2ShortName, "user1", "user2@" BUILTIN_ACL_DOMAIN, false, false, false, true); + writer.Read(topic2ShortName, "user1", "user3@" BUILTIN_ACL_DOMAIN, true, false, false, true); //for topic + writer.Read(topic2ShortName, "user1", "user1@" BUILTIN_ACL_DOMAIN, true, false, false, true); //for consumer + + writer.Read(topic2ShortName, "user2", ticket1, true, false, false, true); + writer.Read(topic2ShortName, "user2", ticket2, false, false, false, true); + + writer.Read(topic2ShortName, "user5", ticket1, true, false, false, true); + writer.Read(topic2ShortName, "user5", ticket2, true, false, false, true); + + acl.Clear(); + acl.AddAccess(NACLib::EAccessType::Allow, NACLib::SelectRow, "user3@" BUILTIN_ACL_DOMAIN); + server.AnnoyingClient->ModifyACL("/Root/PQ", topic2, acl.SerializeAsString()); + + Sleep(TDuration::Seconds(2)); + writer.Read(topic2ShortName, "user1", "user3@" BUILTIN_ACL_DOMAIN, false, true, false, true); + } + + + + Y_UNIT_TEST(CheckKillBalancer) { + TTestServer server; + server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY, NKikimrServices::PQ_READ_PROXY}); + server.EnableLogs({ NKikimrServices::PERSQUEUE}, NActors::NLog::PRI_INFO); + + PrepareForGrpc(server); + + auto writer = MakeDataWriter(server, "source1"); + TTestPQLib PQLib(server); + auto [consumer, p] = PQLib.CreateConsumer({SHORT_TOPIC_NAME}, "shared/user", {}, true); + Cerr << p.Response << "\n"; + + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Type == EMT_LOCK); + auto pp = msg.GetValue().ReadyToRead; + pp.SetValue(TLockInfo()); + + msg = consumer->GetNextMessage(); + msg.Wait(); + + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Type == EMT_LOCK); + pp = msg.GetValue().ReadyToRead; + pp.SetValue(TLockInfo()); + + msg = consumer->GetNextMessage(); + + //TODO: make here infly commits - check results + + server.AnnoyingClient->RestartBalancerTablet(server.CleverServer->GetRuntime(), DEFAULT_TOPIC_NAME); + Cerr << "Balancer killed\n"; + + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Type == EMT_RELEASE); + UNIT_ASSERT(!msg.GetValue().Response.GetRelease().GetCanCommit()); + + msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Type == EMT_RELEASE); + UNIT_ASSERT(!msg.GetValue().Response.GetRelease().GetCanCommit()); + + msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Type == EMT_LOCK); + + msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Type == EMT_LOCK); + + + msg = consumer->GetNextMessage(); + bool res = msg.Wait(TDuration::Seconds(1)); + Y_VERIFY(!res); //no read signalled or lock signals + } + + Y_UNIT_TEST(TestWriteStat) { + auto testWriteStat = [](const TString& originallyProvidedConsumerName, + const TString& consumerName, + const TString& consumerPath) { + auto checkCounters = [](auto monPort, const TString& session, + const std::set<std::string>& canonicalSensorNames, + const TString& clientDc, const TString& originDc, + const TString& client, const TString& consumerPath) { + NJson::TJsonValue counters; + if (clientDc.empty() && originDc.empty()) { + counters = GetClientCountersLegacy(monPort, "pqproxy", session, client, consumerPath); + } else { + counters = GetCountersLegacy(monPort, "pqproxy", session, "account/topic1", + clientDc, originDc, client, consumerPath); + } + const auto sensors = counters["sensors"].GetArray(); + std::set<std::string> sensorNames; + std::transform(sensors.begin(), sensors.end(), + std::inserter(sensorNames, sensorNames.begin()), + [](auto& el) { + return el["labels"]["sensor"].GetString(); + }); + auto equal = sensorNames == canonicalSensorNames; + UNIT_ASSERT(equal); + }; + + auto settings = PQSettings(0, 1, true, "10"); + TTestServer server(settings, false); + server.PrepareNetDataFile(FormNetData()); + server.StartServer(); + server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY, NKikimrServices::NET_CLASSIFIER }); + + const auto monPort = TPortManager().GetPort(); + auto Counters = server.CleverServer->GetGRpcServerRootCounters(); + NActors::TMon Monitoring({monPort, "localhost", 3, "root", "localhost", {}, {}, {}}); + Monitoring.RegisterCountersPage("counters", "Counters", Counters); + Monitoring.Start(); + + auto sender = server.CleverServer->GetRuntime()->AllocateEdgeActor(); + + GetClassifierUpdate(*server.CleverServer, sender); //wait for initializing + + server.AnnoyingClient->CreateTopic("rt3.dc1--account--topic1", 10, 10000, 10000, 2000); + + auto writer = MakeDataWriter(server); + TTestPQLib PQLib(server); + auto [producer, res] = PQLib.CreateProducer("account--topic1", "base64:AAAAaaaa____----12", {}, ECodec::RAW); + + Cerr << res.Response << "\n"; + TInstant st(TInstant::Now()); + for (ui32 i = 1; i <= 5; ++i) { + auto f = producer->Write(i, TString(2000, 'a')); + f.Wait(); + Cerr << f.GetValue().Response << "\n"; + if (i == 5) { + UNIT_ASSERT(TInstant::Now() - st > TDuration::Seconds(3)); + UNIT_ASSERT(f.GetValue().Response.GetAck().GetStat().GetPartitionQuotedTimeMs() <= + f.GetValue().Response.GetAck().GetStat().GetTotalTimeInPartitionQueueMs() + 100); + } + } + + checkCounters(server.CleverServer->GetRuntime()->GetMonPort(), + "writeSession", + { + "BytesWrittenOriginal", + "CompactedBytesWrittenOriginal", + "MessagesWrittenOriginal", + "UncompressedBytesWrittenOriginal" + }, + "", "cluster", "", "" + ); + + checkCounters(monPort, + "writeSession", + { + "BytesInflight", + "BytesInflightTotal", + "Errors", + "SessionsActive", + "SessionsCreated", + "WithoutAuth" + }, + "", "cluster", "", "" + ); + { + auto [consumer, res] = PQLib.CreateConsumer({"account/topic1"}, originallyProvidedConsumerName, + {}, false); + Cerr << res.Response << "\n"; + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Type == EMT_DATA); + + checkCounters(monPort, + "readSession", + { + "Commits", + "PartitionsErrors", + "PartitionsInfly", + "PartitionsLocked", + "PartitionsReleased", + "PartitionsToBeLocked", + "PartitionsToBeReleased", + "WaitsForData" + }, + "", "cluster", "", "" + ); + + checkCounters(monPort, + "readSession", + { + "BytesInflight", + "Errors", + "PipeReconnects", + "SessionsActive", + "SessionsCreated", + "PartsPerSession", + "SessionsWithOldBatchingVersion", + "WithoutAuth" + }, + "", "", consumerName, consumerPath + ); + + checkCounters(server.CleverServer->GetRuntime()->GetMonPort(), + "readSession", + { + "BytesReadFromDC" + }, + "Vla", "", "", "" + ); + + checkCounters(server.CleverServer->GetRuntime()->GetMonPort(), + "readSession", + { + "BytesRead", + "MessagesRead" + }, + "", "Dc1", "", "" + ); + + checkCounters(server.CleverServer->GetRuntime()->GetMonPort(), + "readSession", + { + "BytesRead", + "MessagesRead" + }, + "", "cluster", "", "" + ); + + checkCounters(server.CleverServer->GetRuntime()->GetMonPort(), + "readSession", + { + "BytesRead", + "MessagesRead" + }, + "", "cluster", "", "" + ); + + checkCounters(server.CleverServer->GetRuntime()->GetMonPort(), + "readSession", + { + "BytesRead", + "MessagesRead" + }, + "", "Dc1", consumerName, consumerPath + ); + } + }; + + testWriteStat("some@random@consumer", "some@random@consumer", "some/random/consumer"); + testWriteStat("some@user", "some@user", "some/user"); + testWriteStat("shared@user", "shared@user", "shared/user"); + testWriteStat("shared/user", "user", "shared/user"); + testWriteStat("user", "user", "shared/user"); + testWriteStat("some@random/consumer", "some@random@consumer", "some/random/consumer"); + testWriteStat("/some/user", "some@user", "some/user"); + } + + + Y_UNIT_TEST(TestWriteStat1stClass) { + auto testWriteStat1stClass = [](const TString& consumerPath) { + const auto folderId = "somefolder"; + const auto cloudId = "somecloud"; + const auto databaseId = "root"; + const TString fullTopicName{"/Root/account/topic1"}; + const TString topicName{"account/topic1"}; + + auto checkCounters = + [cloudId, folderId, databaseId](auto monPort, + const std::set<std::string>& canonicalSensorNames, + const TString& stream, const TString& consumer, + const TString& host, const TString& shard) { + auto counters = GetCounters1stClass(monPort, "datastreams", cloudId, databaseId, + folderId, stream, consumer, host, shard); + const auto sensors = counters["sensors"].GetArray(); + std::set<std::string> sensorNames; + std::transform(sensors.begin(), sensors.end(), + std::inserter(sensorNames, sensorNames.begin()), + [](auto& el) { + return el["labels"]["name"].GetString(); + }); + auto equal = sensorNames == canonicalSensorNames; + UNIT_ASSERT(equal); + }; + + auto settings = PQSettings(0, 1, true, "10"); + TTestServer server(settings, false); + server.PrepareNetDataFile(FormNetData()); + server.ServerSettings.PQConfig.SetTopicsAreFirstClassCitizen(true); + server.StartServer(); + server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY, NKikimrServices::NET_CLASSIFIER }); + + UNIT_ASSERT_VALUES_EQUAL(NMsgBusProxy::MSTATUS_OK, + server.AnnoyingClient->AlterUserAttributes("/", "Root", + {{"folder_id", folderId}, + {"cloud_id", cloudId}, + {"database_id", databaseId}})); + + const auto monPort = TPortManager().GetPort(); + auto Counters = server.CleverServer->GetGRpcServerRootCounters(); + NActors::TMon Monitoring({monPort, "localhost", 3, "root", "localhost", {}, {}, {}}); + Monitoring.RegisterCountersPage("counters", "Counters", Counters); + Monitoring.Start(); + + auto sender = server.CleverServer->GetRuntime()->AllocateEdgeActor(); + + GetClassifierUpdate(*server.CleverServer, sender); //wait for initializing + + server.AnnoyingClient->SetNoConfigMode(); + server.AnnoyingClient->FullInit(); + server.AnnoyingClient->InitUserRegistry(); + server.AnnoyingClient->MkDir("/Root", "account"); + server.AnnoyingClient->CreateTopicNoLegacy(fullTopicName, 5); + + NYdb::TDriverConfig driverCfg; + driverCfg.SetEndpoint(TStringBuilder() << "localhost:" << server.GrpcPort) + .SetLog(CreateLogBackend("cerr", ELogPriority::TLOG_DEBUG)).SetDatabase("/Root"); + + auto ydbDriver = MakeHolder<NYdb::TDriver>(driverCfg); + auto persQueueClient = MakeHolder<NYdb::NPersQueue::TPersQueueClient>(*ydbDriver); + + { + auto res = persQueueClient->AddReadRule(fullTopicName, + NYdb::NPersQueue::TAddReadRuleSettings().ReadRule(NYdb::NPersQueue::TReadRuleSettings().ConsumerName(consumerPath))); + res.Wait(); + UNIT_ASSERT(res.GetValue().IsSuccess()); + } + + auto writer = MakeDataWriter(server); + TTestPQLib PQLib(server); + auto [producer, res] = PQLib.CreateProducer(fullTopicName, "base64:AAAAaaaa____----12", {}, + ECodec::RAW); + + Cerr << res.Response << "\n"; + for (ui32 i = 1; i <= 5; ++i) { + auto f = producer->Write(i, TString(2000, 'a')); + f.Wait(); + Cerr << f.GetValue().Response << "\n"; + + } + + { + auto [consumer, res] = PQLib.CreateConsumer({topicName}, consumerPath, {}, + false); + Cerr << res.Response << "\n"; + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + + checkCounters(monPort, + { + "stream.internal_read.commits_per_second", + "stream.internal_read.partitions_errors_per_second", + "stream.internal_read.partitions_locked", + "stream.internal_read.partitions_locked_per_second", + "stream.internal_read.partitions_released_per_second", + "stream.internal_read.partitions_to_be_locked", + "stream.internal_read.partitions_to_be_released", + "stream.internal_read.waits_for_data", + "stream.internal_write.bytes_proceeding", + "stream.internal_write.bytes_proceeding_total", + "stream.internal_write.errors_per_second", + "stream.internal_write.sessions_active", + "stream.internal_write.sessions_created_per_second", + "stream.internal_write.sessions_without_auth" + }, + topicName, "", "", "" + ); + + checkCounters(monPort, + { + "stream.internal_read.commits_per_second", + "stream.internal_read.partitions_errors_per_second", + "stream.internal_read.partitions_locked", + "stream.internal_read.partitions_locked_per_second", + "stream.internal_read.partitions_released_per_second", + "stream.internal_read.partitions_to_be_locked", + "stream.internal_read.partitions_to_be_released", + "stream.internal_read.waits_for_data", + }, + topicName, consumerPath, "", "" + ); + + checkCounters(server.CleverServer->GetRuntime()->GetMonPort(), + { + "stream.internal_read.time_lags_milliseconds", + "stream.incoming_bytes_per_second", + "stream.incoming_records_per_second", + "stream.internal_write.bytes_per_second", + "stream.internal_write.compacted_bytes_per_second", + "stream.internal_write.partition_write_quota_wait_milliseconds", + "stream.internal_write.record_size_bytes", + "stream.internal_write.records_per_second", + "stream.internal_write.time_lags_milliseconds", + "stream.internal_write.uncompressed_bytes_per_second", + "stream.await_operating_milliseconds", + "stream.internal_write.buffer_brimmed_duration_ms", + "stream.internal_read.bytes_per_second", + "stream.internal_read.records_per_second", + "stream.outgoing_bytes_per_second", + "stream.outgoing_records_per_second", + }, + topicName, "", "", "" + ); + + checkCounters(server.CleverServer->GetRuntime()->GetMonPort(), + { + "stream.internal_read.time_lags_milliseconds", + "stream.await_operating_milliseconds", + "stream.internal_read.bytes_per_second", + "stream.internal_read.records_per_second", + "stream.outgoing_bytes_per_second", + "stream.outgoing_records_per_second", + }, + topicName, consumerPath, "", "" + ); + + checkCounters(server.CleverServer->GetRuntime()->GetMonPort(), + { + "stream.await_operating_milliseconds" + }, + topicName, consumerPath, "1", "" + ); + + checkCounters(server.CleverServer->GetRuntime()->GetMonPort(), + { + "stream.internal_write.buffer_brimmed_duration_ms" + }, + topicName, "", "1", "" + ); + } + }; + testWriteStat1stClass("some@random@consumer"); + testWriteStat1stClass("user1"); + } + + + Y_UNIT_TEST(TestUnorderedCommit) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 3); + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + + auto writer = MakeDataWriter(server); + TTestPQLib PQLib(server); + + for (ui32 i = 1; i <= 3; ++i) { + auto [producer, res] = PQLib.CreateProducer(SHORT_TOPIC_NAME, "123" + ToString<ui32>(i), i, ECodec::RAW); + UNIT_ASSERT(res.Response.HasInit()); + auto f = producer->Write(i, TString(10, 'a')); + f.Wait(); + } + auto [consumer, res] = PQLib.CreateConsumer({SHORT_TOPIC_NAME}, "user", 1, false, false); + Cerr << res.Response << "\n"; + for (ui32 i = 1; i <= 3; ++i) { + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasData()); + UNIT_ASSERT(msg.GetValue().Response.GetData().GetCookie() == i); + } + auto msg = consumer->GetNextMessage(); + consumer->Commit({2}); + UNIT_ASSERT(!msg.Wait(TDuration::Seconds(1))); + consumer->Commit({1}); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasCommit()); + UNIT_ASSERT(msg.GetValue().Response.GetCommit().CookieSize() == 2); + UNIT_ASSERT(msg.GetValue().Response.GetCommit().GetCookie(0) == 1); + UNIT_ASSERT(msg.GetValue().Response.GetCommit().GetCookie(1) == 2); + consumer->Commit({3}); + msg = consumer->GetNextMessage(); + UNIT_ASSERT(!msg.Wait(TDuration::MilliSeconds(500))); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasCommit()); + UNIT_ASSERT(msg.GetValue().Response.GetCommit().CookieSize() == 1); + UNIT_ASSERT(msg.GetValue().Response.GetCommit().GetCookie(0) == 3); + consumer->Commit({4}); //not existed cookie + msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasError()); + } + + Y_UNIT_TEST(TestMaxReadCookies) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + auto writer = MakeDataWriter(server); + + TTestPQLib PQLib(server); + { + auto [producer, res] = PQLib.CreateProducer(SHORT_TOPIC_NAME, "123"); + for (ui32 i = 1; i <= 11; ++i) { + auto f = producer->Write(i, TString(10, 'a')); + f.Wait(); + } + } + auto [consumer, res] = PQLib.CreateConsumer({SHORT_TOPIC_NAME}, "user", 1, false); + Cerr << res.Response << "\n"; + for (ui32 i = 1; i <= 11; ++i) { + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasData()); + UNIT_ASSERT(msg.GetValue().Response.GetData().GetCookie() == i); + } + for (ui32 i = 11; i >= 1; --i) { + consumer->Commit({i}); + if (i == 5) { + server.AnnoyingClient->GetClientInfo({DEFAULT_TOPIC_NAME}, "user", true); + } + } + Cerr << "cookies committed\n"; + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasError()); + } + + Y_UNIT_TEST(TestWriteSessionsConflicts) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + auto writer = MakeDataWriter(server); + server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY }); + + TTestPQLib PQLib(server); + auto [producer, res] = PQLib.CreateProducer(SHORT_TOPIC_NAME, "123"); + auto dead = producer->IsDead(); + + auto [producer2, res2] = PQLib.CreateProducer(SHORT_TOPIC_NAME, "123"); + + UNIT_ASSERT(dead.Wait(TDuration::Seconds(1))); + UNIT_ASSERT(!producer2->IsDead().Wait(TDuration::Seconds(1))); + } + + Y_UNIT_TEST(TestWriteSessionNaming) { + TTestServer server; + server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY }); + + + TTestPQLib PQLib(server); + while (true) { + auto [producer, res] = PQLib.CreateProducer("//", "123"); + Cerr << res.Response << "\n"; + UNIT_ASSERT(res.Response.HasError()); + if (res.Response.GetError().GetCode() != NPersQueue::NErrorCode::INITIALIZING) + break; + Sleep(TDuration::Seconds(1)); + } + } + + + Y_UNIT_TEST(TestRelocks) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + + auto writer = MakeDataWriter(server); + + TIntrusivePtr<TCerrLogger> logger = new TCerrLogger(DEBUG_LOG_LEVEL); + TPQLib PQLib; + + TConsumerSettings ss; + ss.ClientId = "user"; + ss.Server = TServerSetting{"localhost", server.GrpcPort}; + ss.Topics.push_back(SHORT_TOPIC_NAME); + ss.UseLockSession = true; + ss.MaxCount = 1; + + auto consumer = PQLib.CreateConsumer(ss, logger, false); + auto p = consumer->Start(); + for (ui32 i = 0; i < 30; ++i) { + server.CleverServer->GetRuntime()->ResetScheduledCount(); + server.AnnoyingClient->RestartPartitionTablets(server.CleverServer->GetRuntime(), DEFAULT_TOPIC_NAME); + Sleep(TDuration::MilliSeconds(1)); + } + p.Wait(); + Cerr << p.GetValue().Response << "\n"; + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetTopic() == DEFAULT_TOPIC_NAME); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetPartition() == 0); + } + + Y_UNIT_TEST(TestLockErrors) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + + auto writer = MakeDataWriter(server); + + TTestPQLib PQLib(server); + + { + auto [producer, res] = PQLib.CreateProducer(SHORT_TOPIC_NAME, "123"); + for (ui32 i = 1; i <= 11; ++i) { + auto f = producer->Write(i, TString(10, 'a')); + f.Wait(); + } + } + auto createConsumer = [&] () { + auto [consumer, res] = PQLib.CreateConsumer({SHORT_TOPIC_NAME}, "user", 1, true); + return std::move(consumer); + }; + auto consumer = createConsumer(); + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetTopic() == DEFAULT_TOPIC_NAME); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetPartition() == 0); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetReadOffset() == 0); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetEndOffset() == 11); + + auto pp = msg.GetValue().ReadyToRead; + pp.SetValue(TLockInfo{0, 5, false}); + auto future = consumer->IsDead(); + future.Wait(); + Cerr << future.GetValue() << "\n"; + + + consumer = createConsumer(); + msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetTopic() == DEFAULT_TOPIC_NAME); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetPartition() == 0); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetReadOffset() == 0); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetEndOffset() == 11); + + pp = msg.GetValue().ReadyToRead; + pp.SetValue(TLockInfo{12, 12, false}); + future = consumer->IsDead(); + future.Wait(); + Cerr << future.GetValue() << "\n"; + + consumer = createConsumer(); + msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetTopic() == DEFAULT_TOPIC_NAME); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetPartition() == 0); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetReadOffset() == 0); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetEndOffset() == 11); + + pp = msg.GetValue().ReadyToRead; + pp.SetValue(TLockInfo{11, 11, false}); + Sleep(TDuration::Seconds(5)); + + + consumer = createConsumer(); + msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetTopic() == DEFAULT_TOPIC_NAME); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetPartition() == 0); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetReadOffset() == 11); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetEndOffset() == 11); + + pp = msg.GetValue().ReadyToRead; + pp.SetValue(TLockInfo{1, 0, true}); + future = consumer->IsDead(); + future.Wait(); + Cerr << future.GetValue() << "\n"; + + consumer = createConsumer(); + msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetTopic() == DEFAULT_TOPIC_NAME); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetPartition() == 0); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetReadOffset() == 11); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetEndOffset() == 11); + + pp = msg.GetValue().ReadyToRead; + pp.SetValue(TLockInfo{0, 0, false}); + future = consumer->IsDead(); + UNIT_ASSERT(!future.Wait(TDuration::Seconds(5))); + } + + + Y_UNIT_TEST(TestLocalChoose) { + TTestServer server(false); + server.ServerSettings.NodeCount = 3; + server.StartServer(); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + + server.EnableLogs({ NKikimrServices::CHOOSE_PROXY }); + + auto channel = grpc::CreateChannel("localhost:"+ToString(server.GrpcPort), grpc::InsecureChannelCredentials()); + auto stub(NKikimrClient::TGRpcServer::NewStub(channel)); + + auto sender = server.CleverServer->GetRuntime()->AllocateEdgeActor(); + + auto nodeId = server.CleverServer->GetRuntime()->GetNodeId(0); + server.CleverServer->GetRuntime()->Send(new IEventHandle( + MakeGRpcProxyStatusID(nodeId), sender, new TEvGRpcProxyStatus::TEvSetup(true, 0, 0)), 0, false + ); + nodeId = server.CleverServer->GetRuntime()->GetNodeId(1); + server.CleverServer->GetRuntime()->Send( + new IEventHandle( + MakeGRpcProxyStatusID(nodeId), sender, new TEvGRpcProxyStatus::TEvSetup(true, 10000000, 1000000) + ), 1, false + ); + nodeId = server.CleverServer->GetRuntime()->GetNodeId(2); + server.CleverServer->GetRuntime()->Send( + new IEventHandle( + MakeGRpcProxyStatusID(nodeId), sender, new TEvGRpcProxyStatus::TEvSetup(true, 10000000, 1000000) + ), 2, false + ); + + grpc::ClientContext context; + NKikimrClient::TChooseProxyRequest request; + request.SetPreferLocalProxy(true); + NKikimrClient::TResponse response; + auto status = stub->ChooseProxy(&context, request, &response); + UNIT_ASSERT(status.ok()); + Cerr << response << "\n"; + UNIT_ASSERT(response.GetStatus() == NMsgBusProxy::MSTATUS_OK); + UNIT_ASSERT(response.GetProxyCookie() == server.CleverServer->GetRuntime()->GetNodeId(0)); + + grpc::ClientContext context2; + request.SetPreferLocalProxy(false); + NKikimrClient::TResponse response2; + status = stub->ChooseProxy(&context2, request, &response2); + UNIT_ASSERT(status.ok()); + Cerr << response2 << "\n"; + UNIT_ASSERT(response2.GetStatus() == NMsgBusProxy::MSTATUS_OK); + UNIT_ASSERT(response2.GetProxyCookie() > server.CleverServer->GetRuntime()->GetNodeId(0)); + } + + + Y_UNIT_TEST(TestRestartBalancer) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + + auto writer = MakeDataWriter(server); + TTestPQLib PQLib(server); + + { + auto [producer, res] = PQLib.CreateProducer(SHORT_TOPIC_NAME, "123"); + for (ui32 i = 1; i < 2; ++i) { + auto f = producer->Write(i, TString(10, 'a')); + f.Wait(); + } + } + auto [consumer, res] = PQLib.CreateConsumer({SHORT_TOPIC_NAME}, "user", 1, true, {}, {}, 100, {}); + Cerr << res.Response << "\n"; + + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetTopic() == DEFAULT_TOPIC_NAME); + UNIT_ASSERT(msg.GetValue().Response.GetLock().GetPartition() == 0); + + + msg.GetValue().ReadyToRead.SetValue(TLockInfo{0,0,false}); + + msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasData()); +// Sleep(TDuration::MilliSeconds(10)); + server.AnnoyingClient->RestartBalancerTablet(server.CleverServer->GetRuntime(), DEFAULT_TOPIC_NAME); + + msg = consumer->GetNextMessage(); + msg.Wait(); + + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT_C(msg.GetValue().Response.HasRelease(), "Response: " << msg.GetValue().Response); + } + + Y_UNIT_TEST(TestBigMessage) { + TTestServer server(false); + server.GrpcServerOptions.SetMaxMessageSize(130_MB); + server.StartServer(); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY }); + + auto writer = MakeDataWriter(server); + TTestPQLib PQLib(server); + auto [producer, res] = PQLib.CreateProducer(SHORT_TOPIC_NAME, "123"); + auto f = producer->Write(1, TString(63_MB, 'a')); + f.Wait(); + Cerr << f.GetValue().Response << "\n"; + UNIT_ASSERT_C(f.GetValue().Response.HasAck(), f.GetValue().Response); + } + + void TestRereadsWhenDataIsEmptyImpl(bool withWait) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + + auto writer = MakeDataWriter(server); + TTestPQLib PQLib(server); + + // Write nonempty data + NKikimr::NPersQueueTests::TRequestWritePQ writeReq(DEFAULT_TOPIC_NAME, 0, "src", 4); + + auto write = [&](const TString& data, bool empty = false) { + NKikimrPQClient::TDataChunk dataChunk; + dataChunk.SetCreateTime(42); + dataChunk.SetSeqNo(++writeReq.SeqNo); + dataChunk.SetData(data); + if (empty) { + dataChunk.SetChunkType(NKikimrPQClient::TDataChunk::GROW); // this guarantees that data will be threated as empty + } + TString serialized; + UNIT_ASSERT(dataChunk.SerializeToString(&serialized)); + server.AnnoyingClient->WriteToPQ(writeReq, serialized); + }; + write("data1"); + write("data2", true); + if (!withWait) { + write("data3"); + } + + auto [consumer, res] = PQLib.CreateConsumer({SHORT_TOPIC_NAME}, "user", 1, false, {}, {}, 1, 1); + UNIT_ASSERT_C(res.Response.HasInit(), res.Response); + auto msg1 = consumer->GetNextMessage().GetValueSync().Response; + + auto assertHasData = [](const NPersQueue::TReadResponse& msg, const TString& data) { + const auto& d = msg.GetData(); + UNIT_ASSERT_VALUES_EQUAL_C(d.MessageBatchSize(), 1, msg); + UNIT_ASSERT_VALUES_EQUAL_C(d.GetMessageBatch(0).MessageSize(), 1, msg); + UNIT_ASSERT_STRINGS_EQUAL_C(d.GetMessageBatch(0).GetMessage(0).GetData(), data, msg); + }; + UNIT_ASSERT_VALUES_EQUAL_C(msg1.GetData().GetCookie(), 1, msg1); + assertHasData(msg1, "data1"); + + auto resp2Future = consumer->GetNextMessage(); + if (withWait) { + // no data + UNIT_ASSERT(!resp2Future.HasValue()); + UNIT_ASSERT(!resp2Future.HasException()); + + // waits and data doesn't arrive + Sleep(TDuration::MilliSeconds(100)); + UNIT_ASSERT(!resp2Future.HasValue()); + UNIT_ASSERT(!resp2Future.HasException()); + + // write data + write("data3"); + } + const auto& msg2 = resp2Future.GetValueSync().Response; + UNIT_ASSERT_VALUES_EQUAL_C(msg2.GetData().GetCookie(), 2, msg2); + assertHasData(msg2, "data3"); + } + + Y_UNIT_TEST(TestRereadsWhenDataIsEmpty) { + TestRereadsWhenDataIsEmptyImpl(false); + } + + Y_UNIT_TEST(TestRereadsWhenDataIsEmptyWithWait) { + TestRereadsWhenDataIsEmptyImpl(true); + } + + + Y_UNIT_TEST(TestLockAfterDrop) { + TTestServer server; + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + + auto writer = MakeDataWriter(server); + TTestPQLib PQLib(server); + + auto [producer, res] = PQLib.CreateProducer(SHORT_TOPIC_NAME, "123"); + auto f = producer->Write(1, TString(1_KB, 'a')); + f.Wait(); + + auto [consumer, res2] = PQLib.CreateConsumer({SHORT_TOPIC_NAME}, "user", 1, true); + Cerr << res2.Response << "\n"; + auto msg = consumer->GetNextMessage(); + msg.Wait(); + UNIT_ASSERT_C(msg.GetValue().Response.HasLock(), msg.GetValue().Response); + UNIT_ASSERT_C(msg.GetValue().Response.GetLock().GetTopic() == DEFAULT_TOPIC_NAME, msg.GetValue().Response); + UNIT_ASSERT_C(msg.GetValue().Response.GetLock().GetPartition() == 0, msg.GetValue().Response); + + server.CleverServer->GetRuntime()->ResetScheduledCount(); + server.AnnoyingClient->RestartPartitionTablets(server.CleverServer->GetRuntime(), DEFAULT_TOPIC_NAME); + + msg.GetValue().ReadyToRead.SetValue({0,0,false}); + + msg = consumer->GetNextMessage(); + UNIT_ASSERT(msg.Wait(TDuration::Seconds(10))); + + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasData()); + } + + + Y_UNIT_TEST(TestMaxNewTopicModel) { + TTestServer server; + server.AnnoyingClient->AlterUserAttributes("/", "Root", {{"__extra_path_symbols_allowed", "@"}}); + server.AnnoyingClient->CreateTopic("rt3.dc1--aaa@bbb@ccc--topic", 1); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + + auto writer = MakeDataWriter(server); + TTestPQLib PQLib(server); + { + auto [producer, res] = PQLib.CreateProducer("aaa/bbb/ccc/topic", "123"); + UNIT_ASSERT_C(res.Response.HasInit(), res.Response); + for (ui32 i = 1; i <= 11; ++i) { + auto f = producer->Write(i, TString(10, 'a')); + f.Wait(); + UNIT_ASSERT_C(f.GetValue().Response.HasAck(), f.GetValue().Response); + } + } + + auto [consumer, res2] = PQLib.CreateConsumer({"aaa/bbb/ccc/topic"}, "user", 1, true); + UNIT_ASSERT_C(res2.Response.HasInit(), res2.Response); + + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Response.HasLock()); + } + + Y_UNIT_TEST(TestPartitionStatus) { + TTestServer server; + server.AnnoyingClient->AlterUserAttributes("/", "Root", {{"__extra_path_symbols_allowed", "@"}}); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + + auto writer = MakeDataWriter(server); + TTestPQLib PQLib(server); + + auto [consumer, res] = PQLib.CreateConsumer({SHORT_TOPIC_NAME}, "user", 1, true); + UNIT_ASSERT_C(res.Response.HasInit(), res.Response); + + auto msg = consumer->GetNextMessage(); + auto value = msg.ExtractValueSync(); + Cerr << value.Response << "\n"; + UNIT_ASSERT(value.Response.HasLock()); + value.ReadyToRead.SetValue(TLockInfo{}); + auto lock = value.Response.GetLock(); + consumer->RequestPartitionStatus(lock.GetTopic(), lock.GetPartition(), lock.GetGeneration()); + msg = consumer->GetNextMessage(); + value = msg.ExtractValueSync(); + Cerr << value.Response << "\n"; + UNIT_ASSERT(value.Response.HasPartitionStatus()); + UNIT_ASSERT(value.Response.GetPartitionStatus().GetCommittedOffset() == 0); + UNIT_ASSERT(value.Response.GetPartitionStatus().GetEndOffset() == 0); + auto wt = value.Response.GetPartitionStatus().GetWriteWatermarkMs(); + Sleep(TDuration::Seconds(15)); + + consumer->RequestPartitionStatus(lock.GetTopic(), lock.GetPartition(), lock.GetGeneration()); + msg = consumer->GetNextMessage(); + value = msg.ExtractValueSync(); + Cerr << value.Response << "\n"; + UNIT_ASSERT(value.Response.HasPartitionStatus()); + UNIT_ASSERT(value.Response.GetPartitionStatus().GetCommittedOffset() == 0); + UNIT_ASSERT(value.Response.GetPartitionStatus().GetEndOffset() == 0); + UNIT_ASSERT(wt < value.Response.GetPartitionStatus().GetWriteWatermarkMs()); + wt = value.Response.GetPartitionStatus().GetWriteWatermarkMs(); + + { + auto [producer, res] = PQLib.CreateProducer(SHORT_TOPIC_NAME, "123"); + UNIT_ASSERT_C(res.Response.HasInit(), res.Response); + auto f = producer->Write(1, TString(10, 'a')); + UNIT_ASSERT_C(f.GetValueSync().Response.HasAck(), f.GetValueSync().Response); + } + msg = consumer->GetNextMessage(); + value = msg.ExtractValueSync(); + Cerr << value.Response << "\n"; + UNIT_ASSERT(value.Response.HasData()); + auto cookie = value.Response.GetData().GetCookie(); + + consumer->RequestPartitionStatus(lock.GetTopic(), lock.GetPartition(), lock.GetGeneration()); + msg = consumer->GetNextMessage(); + value = msg.ExtractValueSync(); + Cerr << value.Response << "\n"; + UNIT_ASSERT(value.Response.HasPartitionStatus()); + UNIT_ASSERT(value.Response.GetPartitionStatus().GetCommittedOffset() == 0); + UNIT_ASSERT(value.Response.GetPartitionStatus().GetEndOffset() == 1); + UNIT_ASSERT(wt < value.Response.GetPartitionStatus().GetWriteWatermarkMs()); + wt = value.Response.GetPartitionStatus().GetWriteWatermarkMs(); + consumer->Commit({cookie}); + msg = consumer->GetNextMessage(); + Cerr << msg.GetValueSync().Response << "\n"; + UNIT_ASSERT(msg.GetValueSync().Response.HasCommit()); + + consumer->RequestPartitionStatus(lock.GetTopic(), lock.GetPartition(), lock.GetGeneration()); + msg = consumer->GetNextMessage(); + value = msg.ExtractValueSync(); + Cerr << value.Response << "\n"; + UNIT_ASSERT(value.Response.HasPartitionStatus()); + UNIT_ASSERT(value.Response.GetPartitionStatus().GetCommittedOffset() == 1); + UNIT_ASSERT(value.Response.GetPartitionStatus().GetEndOffset() == 1); + } + + Y_UNIT_TEST(TestDeletionOfTopic) { + TTestServer server(false); + server.GrpcServerOptions.SetMaxMessageSize(130_MB); + server.StartServer(); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 1); + server.EnableLogs({ NKikimrServices::PQ_READ_PROXY }); + + auto writer = MakeDataWriter(server); + TTestPQLib PQLib(server); + + server.AnnoyingClient->DescribeTopic({DEFAULT_TOPIC_NAME}); + server.AnnoyingClient->DeleteTopic2(DEFAULT_TOPIC_NAME, NPersQueue::NErrorCode::OK, false); + Sleep(TDuration::Seconds(1)); + auto [consumer, res] = PQLib.CreateConsumer({SHORT_TOPIC_NAME}, "user", 1, {}, {}, false); + Cerr << res.Response << "\n"; + + UNIT_ASSERT_EQUAL_C(res.Response.GetError().GetCode(), NPersQueue::NErrorCode::UNKNOWN_TOPIC, res.Response); + } + + Y_UNIT_TEST(SetupYqlTimeout) { + TTestServer server(PQSettings(0, 1, true, "1")); + server.EnableLogs({ NKikimrServices::PQ_WRITE_PROXY }); + server.AnnoyingClient->CreateTopic(DEFAULT_TOPIC_NAME, 10); + + auto writer = MakeDataWriter(server); + + server.AnnoyingClient->MarkNodeInHive(server.CleverServer->GetRuntime(), 0, false); + server.AnnoyingClient->KickNodeInHive(server.CleverServer->GetRuntime(), 0); + + writer.InitSession("sid1", 2, false); + } + + Y_UNIT_TEST(TestFirstClassWriteAndRead) { + auto settings = PQSettings(0, 1, true, "1"); + settings.PQConfig.SetTopicsAreFirstClassCitizen(true); + TTestServer server(settings, false); + server.StartServer(false); + UNIT_ASSERT_VALUES_EQUAL(NMsgBusProxy::MSTATUS_OK, + server.AnnoyingClient->AlterUserAttributes("/", "Root", {{"folder_id", "somefolder"}, {"cloud_id", "somecloud"}, {"database_id", "root"}})); + + + Cerr << "HERE\n"; + + PrepareForFstClass(*server.AnnoyingClient); + server.EnableLogs({ NKikimrServices::PQ_METACACHE, NKikimrServices::PERSQUEUE }, NActors::NLog::PRI_INFO); + server.EnableLogs({ NKikimrServices::PERSQUEUE_CLUSTER_TRACKER }, NActors::NLog::PRI_EMERG); + server.EnableLogs({ NKikimrServices::PERSQUEUE_READ_BALANCER }, NActors::NLog::PRI_DEBUG); + + TTestPQLib pqLib(server); + + NACLib::TDiffACL acl; +// acl.AddAccess(NACLib::EAccessType::Allow, NACLib::UpdateRow, "topic@" BUILTIN_ACL_DOMAIN); +// server.AnnoyingClient->ModifyACL("/Root/account1", "root-acc-topic", acl.SerializeAsString()); + Sleep(TDuration::Seconds(5)); + { + auto [producer, res] = pqLib.CreateProducer(FC_TOPIC_PATH, "123", {}, {}, {}, true); + UNIT_ASSERT_C(res.Response.HasInit(), res.Response); + Cerr << "Producer 1 start response: " << res.Response << "\n"; + auto f = producer->Write(1, TString(10, 'a')); + UNIT_ASSERT_C(f.GetValueSync().Response.HasAck(), f.GetValueSync().Response); + } + { + auto [producer, res] = pqLib.CreateProducer(FC_TOPIC_PATH, "123", {}, {}, {}, true); + UNIT_ASSERT_C(res.Response.HasInit(), res.Response); + Cerr << "Producer 2 start response: " << res.Response << "\n"; + auto f = producer->Write(2, TString(10, 'b')); + UNIT_ASSERT_C(f.GetValueSync().Response.HasAck(), f.GetValueSync().Response); + } + { + auto [consumer, res] = pqLib.CreateConsumer({FC_TOPIC_PATH}, "user"); + Cerr << "Consumer start response: " << res.Response << "\n"; + auto msg = consumer->GetNextMessage(); + msg.Wait(); + Cerr << "Read response: " << msg.GetValue().Response << "\n"; + UNIT_ASSERT(msg.GetValue().Type == EMT_DATA); + } + { + auto [consumer, res] = pqLib.CreateConsumer({FC_TOPIC_PATH}, "user", {}, true); + Cerr << "Consumer start response: " << res.Response << "\n"; + ui64 commitsDone = 0; + while (true) { + auto msg = consumer->GetNextMessage(); + msg.Wait(); + auto& value = msg.GetValue(); + if (value.Type == EMT_LOCK) { + TStringBuilder message; + Cerr << "Consumer lock response: " << value.Response << "\n"; + UNIT_ASSERT_VALUES_EQUAL(value.Response.GetLock().GetTopic(), "account1/root-acc-topic"); + msg.GetValue().ReadyToRead.SetValue(TLockInfo()); + } else if (value.Type == EMT_DATA) { + auto cookie = msg.GetValue().Response.GetData().GetCookie(); + consumer->Commit({cookie}); + } else { + UNIT_ASSERT(value.Type == EMT_COMMIT); + commitsDone++; + break; + } + } + UNIT_ASSERT(commitsDone > 0); + } + } + + Y_UNIT_TEST(SrcIdCompatibility) { + TString srcId1 = "test-src-id-compat", srcId2 = "test-src-id-compat2"; + TTestServer server{}; + TString topicName = "rt3.dc1--topic100"; + TString fullPath = "Root/PQ/rt3.dc1--topic100"; + TString shortTopicName = "topic100"; + server.AnnoyingClient->CreateTopic(topicName, 100); + server.EnableLogs({ NKikimrServices::PERSQUEUE }, NActors::NLog::PRI_INFO); + + auto runTest = [&] ( + const TString& topicToAdd, const TString& topicForHash, const TString& topicName, + const TString& srcId, ui32 partId, ui64 accessTime = 0 + ) { + TStringBuilder query; + auto encoded = NPQ::NSourceIdEncoding::EncodeSrcId(topicForHash, srcId); + Cerr << "===save partition with time: " << accessTime << Endl; + + if (accessTime == 0) { + accessTime = TInstant::Now().MilliSeconds(); + } + if (!topicToAdd.empty()) { // Empty means don't add anything + query << + "--!syntax_v1\n" + "UPSERT INTO `/Root/PQ/SourceIdMeta2` (Hash, Topic, SourceId, CreateTime, AccessTime, Partition) VALUES (" + << encoded.Hash << ", \"" << topicToAdd << "\", \"" << encoded.EscapedSourceId << "\", " + << TInstant::Now().MilliSeconds() << ", " << accessTime << ", " << partId << "); "; + Cerr << "Run query:\n" << query << Endl; + auto scResult = server.AnnoyingClient->RunYqlDataQuery(query); + //UNIT_ASSERT(scResult.Defined()); + } + TTestPQLib pqLib(server); + auto[producer, response] = pqLib.CreateProducer(topicName, srcId, {}, {}, {}, true); + UNIT_ASSERT_C(response.Response.HasInit(), response.Response); + UNIT_ASSERT_VALUES_EQUAL(response.Response.GetInit().GetPartition(), partId); + }; + + runTest(fullPath, shortTopicName, shortTopicName, srcId1, 5, 100); + runTest(topicName, shortTopicName, shortTopicName, srcId2, 6); + runTest("", "", shortTopicName, srcId1, 5, 100); + // Add another partition to the src mapping with different topic in key. + // Expect newer partition to be discovered. + ui64 time = (TInstant::Now() + TDuration::Hours(4)).MilliSeconds(); + runTest(topicName, shortTopicName, shortTopicName, srcId1, 7, time); + + } + +} // Y_UNIT_TEST_SUITE(TPersQueueTest) + +} // namespace NPersQueueTests +} // namespace NKikimr diff --git a/kikimr/yndx/grpc_services/persqueue/protocol_compatibility_ut.cpp b/kikimr/yndx/grpc_services/persqueue/protocol_compatibility_ut.cpp new file mode 100644 index 0000000000..f69080ea6b --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/protocol_compatibility_ut.cpp @@ -0,0 +1,80 @@ +#include <ydb/public/sdk/cpp/client/ydb_persqueue_core/ut/ut_utils/test_server.h>
+#include <ydb/core/client/server/msgbus_server_pq_metacache.h>
+#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h>
+#include <kikimr/yndx/grpc_services/persqueue/persqueue.h>
+#include <kikimr/yndx/persqueue/msgbus_server/read_session_info.h>
+
+
+namespace NKikimr {
+namespace NPersQueueTests {
+
+Y_UNIT_TEST_SUITE(TPersQueueProtocolCompatibility) {
+ Y_UNIT_TEST(GetReadInfoFromV1AboutV0Session) {
+ NKikimr::Tests::TServerSettings serverSettings = PQSettings(0);
+ serverSettings.RegisterGrpcService<NKikimr::NGRpcService::TGRpcPersQueueService>(
+ "pq",
+ NKikimr::NMsgBusProxy::CreatePersQueueMetaCacheV2Id()
+ );
+ serverSettings.SetPersQueueGetReadSessionsInfoWorkerFactory(
+ std::make_shared<NKikimr::NMsgBusProxy::TPersQueueGetReadSessionsInfoWorkerWithPQv0Factory>()
+ );
+
+ NPersQueue::TTestServer server(serverSettings);
+ server.EnableLogs({ NKikimrServices::PERSQUEUE, NKikimrServices::PQ_READ_PROXY });
+ server.AnnoyingClient->CreateTopic("rt3.dc1--topic1", 1);
+
+ NPersQueue::TPQLibSettings pqSettings;
+ pqSettings.DefaultLogger = new NPersQueue::TCerrLogger(NPersQueue::DEBUG_LOG_LEVEL);
+ THolder<NPersQueue::TPQLib> PQLib = MakeHolder<NPersQueue::TPQLib>(pqSettings);
+
+ NPersQueue::TConsumerSettings settings;
+ settings.Server = NPersQueue::TServerSetting{"localhost", server.GrpcPort};
+ settings.ClientId = "user";
+ settings.Topics = {"topic1"};
+ settings.UseLockSession = true;
+ auto consumer = PQLib->CreateConsumer(settings);
+ auto response = consumer->Start().GetValueSync();
+ UNIT_ASSERT_C(response.Response.HasInit(), response.Response);
+
+ auto msg = consumer->GetNextMessage();
+ auto value = msg.ExtractValueSync();
+ Cerr << value.Response << "\n";
+ UNIT_ASSERT(value.Response.HasLock());
+ value.ReadyToRead.SetValue(NPersQueue::TLockInfo{});
+ auto lock = value.Response.GetLock();
+ Cout << lock.DebugString() << Endl;
+ {
+ std::shared_ptr<grpc::Channel> channel;
+ std::unique_ptr<Ydb::PersQueue::V1::PersQueueService::Stub> stub;
+
+ {
+ channel = grpc::CreateChannel(
+ "localhost:" + ToString(server.GrpcPort),
+ grpc::InsecureChannelCredentials()
+ );
+ stub = Ydb::PersQueue::V1::PersQueueService::NewStub(channel);
+ }
+ {
+ Sleep(TDuration::Seconds(10));
+ Ydb::PersQueue::V1::ReadInfoRequest request;
+ Ydb::PersQueue::V1::ReadInfoResponse response;
+ request.mutable_consumer()->set_path("user");
+ request.set_get_only_original(true);
+ request.add_topics()->set_path("topic1");
+ grpc::ClientContext rcontext;
+ auto status = stub->GetReadSessionsInfo(&rcontext, request, &response);
+ UNIT_ASSERT(status.ok());
+ Ydb::PersQueue::V1::ReadInfoResult res;
+ response.operation().result().UnpackTo(&res);
+ Cerr << "Read info response: " << response << Endl << res << Endl;
+ UNIT_ASSERT_VALUES_EQUAL(res.topics_size(), 1);
+ UNIT_ASSERT(res.topics(0).status() == Ydb::StatusIds::SUCCESS);
+ }
+ }
+
+ }
+
+} // Y_UNIT_TEST_SUITE(TPersQueueProtocolCompatibility)
+
+} // namespace NPersQueueTests
+} // namespace NKikimr
diff --git a/kikimr/yndx/grpc_services/persqueue/ut/definitions.h b/kikimr/yndx/grpc_services/persqueue/ut/definitions.h new file mode 100644 index 0000000000..35f9f6bc43 --- /dev/null +++ b/kikimr/yndx/grpc_services/persqueue/ut/definitions.h @@ -0,0 +1,18 @@ +#pragma once +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/data_writer.h> +#include <kikimr/persqueue/sdk/deprecated/cpp/v2/ut_utils/test_server.h> + +namespace NKikimr::NPersQueueTests { + + +namespace { + const TString DEFAULT_TOPIC_NAME = "rt3.dc1--topic1"; + const TString FC_TOPIC_PATH = "/Root/account1/root-acc-topic"; + const TString SHORT_TOPIC_NAME = "topic1"; +} + +NPersQueue::NTests::TPQDataWriter MakeDataWriter(NPersQueue::TTestServer& server, const TString& srcId = "source") { + return NPersQueue::NTests::TPQDataWriter(DEFAULT_TOPIC_NAME, SHORT_TOPIC_NAME, srcId, server); +} + +} // namespace NKikimr::NPersQueueTests diff --git a/kikimr/yndx/persqueue/msgbus_server/CMakeLists.txt b/kikimr/yndx/persqueue/msgbus_server/CMakeLists.txt new file mode 100644 index 0000000000..c923814a2f --- /dev/null +++ b/kikimr/yndx/persqueue/msgbus_server/CMakeLists.txt @@ -0,0 +1,19 @@ + +# This file was gererated by the build system used internally in the Yandex monorepo. +# Only simple modifications are allowed (adding source-files to targets, adding simple properties +# like target_include_directories). These modifications will be ported to original +# ya.make files by maintainers. Any complex modifications which can't be ported back to the +# original buildsystem will not be accepted. + + + +add_library(yndx-persqueue-msgbus_server) +target_link_libraries(yndx-persqueue-msgbus_server PUBLIC + contrib-libs-cxxsupp + yutil + core-client-server + yndx-grpc_services-persqueue +) +target_sources(yndx-persqueue-msgbus_server PRIVATE + ${CMAKE_SOURCE_DIR}/kikimr/yndx/persqueue/msgbus_server/read_session_info.cpp +) diff --git a/kikimr/yndx/persqueue/msgbus_server/read_session_info.cpp b/kikimr/yndx/persqueue/msgbus_server/read_session_info.cpp new file mode 100644 index 0000000000..1ccf3b0539 --- /dev/null +++ b/kikimr/yndx/persqueue/msgbus_server/read_session_info.cpp @@ -0,0 +1,16 @@ +#include "read_session_info.h" + + +namespace NKikimr { +namespace NMsgBusProxy { + +void TPersQueueGetReadSessionsInfoWorkerWithPQv0::SendStatusRequest(const TString& sessionName, TActorId actorId, const TActorContext& ctx) { + if (sessionName.EndsWith("_v1")) { + SendStatusRequest<NGRpcProxy::V1::TEvPQProxy::TEvReadSessionStatus>(actorId, ctx); + } else { + SendStatusRequest<NGRpcProxy::TEvPQProxy::TEvReadSessionStatus>(actorId, ctx); + } +} + +} // namespace NMsgBusProxy +} // namespace NKikimr diff --git a/kikimr/yndx/persqueue/msgbus_server/read_session_info.h b/kikimr/yndx/persqueue/msgbus_server/read_session_info.h new file mode 100644 index 0000000000..f92572677d --- /dev/null +++ b/kikimr/yndx/persqueue/msgbus_server/read_session_info.h @@ -0,0 +1,43 @@ +#pragma once + +#include <ydb/core/client/server/msgbus_server_pq_read_session_info.h> + +#include <kikimr/yndx/grpc_services/persqueue/grpc_pq_actor.h> + + +namespace NKikimr { +namespace NMsgBusProxy { + +class TPersQueueGetReadSessionsInfoWorkerWithPQv0 : public IPersQueueGetReadSessionsInfoWorker { +public: + using TBase = IPersQueueGetReadSessionsInfoWorker; + using TBase::TBase; + using TBase::SendStatusRequest; + + STFUNC(StateFunc) override { + switch (ev->GetTypeRewrite()) { + HFunc(NGRpcProxy::TEvPQProxy::TEvReadSessionStatusResponse, HandleStatusResponse<NGRpcProxy::TEvPQProxy::TEvReadSessionStatusResponse>); + HFunc(NGRpcProxy::V1::TEvPQProxy::TEvReadSessionStatusResponse, HandleStatusResponse<NGRpcProxy::V1::TEvPQProxy::TEvReadSessionStatusResponse>); + HFunc(TEvents::TEvUndelivered, Undelivered); + HFunc(TEvInterconnect::TEvNodeDisconnected, Disconnected); + } + } + +private: + void SendStatusRequest(const TString& sessionName, TActorId actorId, const TActorContext& ctx) override; +}; + +class TPersQueueGetReadSessionsInfoWorkerWithPQv0Factory : public IPersQueueGetReadSessionsInfoWorkerFactory { +public: + THolder<IPersQueueGetReadSessionsInfoWorker> Create( + const TActorId& parentId, + const THashMap<TString, TActorId>& readSessions, + std::shared_ptr<const TPersQueueBaseRequestProcessor::TNodesInfo> nodesInfo + ) const override { + return MakeHolder<TPersQueueGetReadSessionsInfoWorkerWithPQv0>(parentId, readSessions, nodesInfo); + } +}; + +} // namespace NMsgBusProxy +} // namespace NKikimr + diff --git a/kikimr/yndx/persqueue/read_batch_converter/CMakeLists.txt b/kikimr/yndx/persqueue/read_batch_converter/CMakeLists.txt new file mode 100644 index 0000000000..9e8f9b6006 --- /dev/null +++ b/kikimr/yndx/persqueue/read_batch_converter/CMakeLists.txt @@ -0,0 +1,18 @@ + +# This file was gererated by the build system used internally in the Yandex monorepo. +# Only simple modifications are allowed (adding source-files to targets, adding simple properties +# like target_include_directories). These modifications will be ported to original +# ya.make files by maintainers. Any complex modifications which can't be ported back to the +# original buildsystem will not be accepted. + + + +add_library(yndx-persqueue-read_batch_converter) +target_link_libraries(yndx-persqueue-read_batch_converter PUBLIC + contrib-libs-cxxsupp + yutil + api-protos-yndx +) +target_sources(yndx-persqueue-read_batch_converter PRIVATE + ${CMAKE_SOURCE_DIR}/kikimr/yndx/persqueue/read_batch_converter/read_batch_converter.cpp +) diff --git a/kikimr/yndx/persqueue/read_batch_converter/read_batch_converter.cpp b/kikimr/yndx/persqueue/read_batch_converter/read_batch_converter.cpp new file mode 100644 index 0000000000..bca03dc72f --- /dev/null +++ b/kikimr/yndx/persqueue/read_batch_converter/read_batch_converter.cpp @@ -0,0 +1,43 @@ +#include "read_batch_converter.h" + +namespace NPersQueue { + +static void Convert(const ReadResponse::BatchedData::PartitionData& partition, ReadResponse::Data::MessageBatch* dstBatch) { + dstBatch->set_topic(partition.topic()); + dstBatch->set_partition(partition.partition()); + for (const ReadResponse::BatchedData::Batch& batch : partition.batch()) { + for (const ReadResponse::BatchedData::MessageData& message : batch.message_data()) { + ReadResponse::Data::Message* const dstMessage = dstBatch->add_message(); + dstMessage->set_data(message.data()); + dstMessage->set_offset(message.offset()); + + MessageMeta* const meta = dstMessage->mutable_meta(); + meta->set_source_id(batch.source_id()); + meta->set_seq_no(message.seq_no()); + meta->set_create_time_ms(message.create_time_ms()); + meta->set_write_time_ms(batch.write_time_ms()); + meta->set_codec(message.codec()); + meta->set_ip(batch.ip()); + meta->set_uncompressed_size(message.uncompressed_size()); + if (batch.has_extra_fields()) { + *meta->mutable_extra_fields() = batch.extra_fields(); + } + } + } +} + +void ConvertToOldBatch(ReadResponse& response) { + if (!response.has_batched_data()) { + return; + } + ReadResponse::BatchedData data; + data.Swap(response.mutable_batched_data()); + + ReadResponse::Data& dstData = *response.mutable_data(); // this call will clear BatchedData field + dstData.set_cookie(data.cookie()); + for (const ReadResponse::BatchedData::PartitionData& partition : data.partition_data()) { + Convert(partition, dstData.add_message_batch()); + } +} + +} // namespace NPersQueue diff --git a/kikimr/yndx/persqueue/read_batch_converter/read_batch_converter.h b/kikimr/yndx/persqueue/read_batch_converter/read_batch_converter.h new file mode 100644 index 0000000000..ddf32da4ae --- /dev/null +++ b/kikimr/yndx/persqueue/read_batch_converter/read_batch_converter.h @@ -0,0 +1,10 @@ +#pragma once +#include <kikimr/yndx/api/protos/persqueue.pb.h> + +namespace NPersQueue { + +// Converts responses with BatchedData field to responses with Data field. +// Other responses will be leaved unchanged. +void ConvertToOldBatch(ReadResponse& response); + +} // namespace NPersQueue |