diff --git a/.dockerignore b/.dockerignore deleted file mode 100644 index ea681c12685..00000000000 --- a/.dockerignore +++ /dev/null @@ -1,47 +0,0 @@ -# Git and version control -.git -.gitignore -.github -.gitattributes - -# IDE and editor files -.idea/ -.vscode/ -.cache/ -.cursor/ -.cursorrules -*.swp -*.swo -.DS_Store -*.dSYM -.clangd -compile_commands.json - -# Build artifacts and cache -*.o -*.a -*.so -*.dylib -*.exe -*.out -*.log -*.pyc -__pycache__/ -*.egg-info/ -.pytest_cache/ -node_modules/ - -# Documentation -*.md -docs/ - -# CI/CD -.circleci -.travis.yml -azure-pipelines.yml - -# Large directories not needed for SDK builds -ydb/core/ -ydb/apps/ -ydb/services/ -ydb/tools/ diff --git a/.github/import_generation.txt b/.github/import_generation.txt index 920a1396648..c739b42c4d2 100644 --- a/.github/import_generation.txt +++ b/.github/import_generation.txt @@ -1 +1 @@ -43 +44 diff --git a/.github/last_commit.txt b/.github/last_commit.txt index 7fd1813ea92..3a09c931bbf 100644 --- a/.github/last_commit.txt +++ b/.github/last_commit.txt @@ -1 +1 @@ -b3b57fcb5915100b21bf9a507bf0a8c0627b7fee +fb25fb6453264260a02c6ec49492382fcaf801ce diff --git a/.github/scripts/run_iam_integration_tests.sh b/.github/scripts/run_iam_integration_tests.sh new file mode 100755 index 00000000000..43cf89f56c8 --- /dev/null +++ b/.github/scripts/run_iam_integration_tests.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash + +set -euo pipefail + +IAM_REGEX='^(DriverAuth|TMetadataFixture|TJwtIamFixture|TOAuthIamFixture|OAuth_WithFacility)\.' +IAM_CONTAINER_NAME="${IAM_CONTAINER_NAME:-ydb-iam}" +IAM_CTEST_JOBS="${IAM_CTEST_JOBS:-2}" +IAM_READY_ATTEMPTS="${IAM_READY_ATTEMPTS:-60}" +IAM_READY_SLEEP_SECONDS="${IAM_READY_SLEEP_SECONDS:-2}" +IAM_TOKEN="${IAM_TOKEN:-root@builtin}" + +cleanup_iam() { + docker rm -f "${IAM_CONTAINER_NAME}" >/dev/null 2>&1 || true +} + +wait_for_iam_ydb() { + for _ in $(seq 1 "${IAM_READY_ATTEMPTS}"); do + if docker exec -e "YDB_TOKEN=${IAM_TOKEN}" "${IAM_CONTAINER_NAME}" /ydb \ + --endpoint grpc://localhost:2136 \ + --database /local \ + sql -s 'select 1' >/dev/null 2>&1; then + return 0 + fi + sleep "${IAM_READY_SLEEP_SECONDS}" + done + + echo "ydb-iam failed to become ready" + docker logs "${IAM_CONTAINER_NAME}" 2>&1 | tail -50 || true + return 1 +} + +trap cleanup_iam EXIT +cleanup_iam + +docker run -d --name "${IAM_CONTAINER_NAME}" --hostname localhost \ + -p 2235:2135 -p 2236:2136 -p 28765:8765 \ + -v /tmp/ydb_iam_certs:/ydb_certs \ + -e YDB_USE_IN_MEMORY_PDISKS=true \ + -e YDB_TABLE_ENABLE_PREPARED_DDL=true \ + -e YDB_ENFORCE_USER_TOKEN_REQUIREMENT=true \ + -e YDB_DEFAULT_CLUSTERADMIN=root@builtin \ + ghcr.io/ydb-platform/local-ydb:trunk + +wait_for_iam_ydb + +YDB_ENDPOINT=localhost:2236 YDB_DATABASE=/local \ + ctest -j"${IAM_CTEST_JOBS}" --test-dir build -R "${IAM_REGEX}" --output-on-failure diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 7baf56e7414..caeda7f23f7 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -61,7 +61,15 @@ jobs: - name: Test shell: bash run: | - ctest -j1 --preset coverage-all --output-on-failure + set -euo pipefail + + IAM_REGEX='^(DriverAuth|TMetadataFixture|TJwtIamFixture|TOAuthIamFixture|OAuth_WithFacility)\.' + FLAKY_REGEX='(ManyMessages|DiscoveryHang|DescribeHang)' + EXCLUDE_REGEX="${IAM_REGEX}|${FLAKY_REGEX}" + + ctest -j1 --preset coverage-all -E "${EXCLUDE_REGEX}" --output-on-failure + + IAM_CTEST_JOBS=1 ./.github/scripts/run_iam_integration_tests.sh - name: Generate coverage report uses: ./.github/actions/gcovr diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index f0816775abb..bf88544d103 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -117,7 +117,16 @@ jobs: - name: Test shell: bash run: | - YDB_VERSION=${{ matrix.ydb-version }} ctest -j2 --preset integration + IAM_REGEX='^(DriverAuth|TMetadataFixture|TJwtIamFixture|TOAuthIamFixture|OAuth_WithFacility)\.' + + YDB_VERSION=${{ matrix.ydb-version }} ctest -j2 --preset integration \ + -E "${IAM_REGEX}" --output-on-failure + + case '${{ matrix.ydb-version }}' in + 25.1|trunk) + ./.github/scripts/run_iam_integration_tests.sh + ;; + esac test-install: name: "Test CMake Install" @@ -267,4 +276,4 @@ jobs: - name: Smoke-test generated .deb packages shell: bash run: | - ./scripts/test_deb_packages.sh build-deb \ No newline at end of file + ./scripts/test_deb_packages.sh build-deb diff --git a/CHANGELOG.md b/CHANGELOG.md index 80e8d11280b..f9d4fe04a0c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,9 @@ +## v3.20.0 + +* Added automatic retries for unary methods of table and query clients(ExecuteQuery, ExecuteScript, BulkUpsert, ReadRows). + +* Implemented native ranges(TRowRange) and iterators over both streaming query results and TResultSet. + ## v3.19.0 * Added metric buffer for batched metric writes to reduce OpenTelemetry hot-path overhead. diff --git a/Testing/Temporary/LastTest.log b/Testing/Temporary/LastTest.log new file mode 100644 index 00000000000..7d298711f79 --- /dev/null +++ b/Testing/Temporary/LastTest.log @@ -0,0 +1,3 @@ +Start testing: Jul 01 00:28 MSK +---------------------------------------------------------- +End testing: Jul 01 00:28 MSK diff --git a/examples/basic_example/basic_example.cpp b/examples/basic_example/basic_example.cpp index cc5839b0475..f44a58d65b9 100644 --- a/examples/basic_example/basic_example.cpp +++ b/examples/basic_example/basic_example.cpp @@ -1,14 +1,14 @@ #include "basic_example.h" #include - +#include #include -#include using namespace NYdb; using namespace NYdb::NQuery; using namespace NYdb::NStatusHelpers; +using NYdb::NRetry::TRetryOperationSettings; template std::string OptionalToString(const std::optional& opt) { @@ -338,14 +338,18 @@ void MultiStep(TQueryClient client) { return result2; })); // The end of the retried lambda - TResultSetParser parser(*resultSet); std::cout << "> MultiStep:" << std::endl; - while (parser.TryNextRow()) { - auto airDate = TInstant::Days(*parser.ColumnParser("air_date").GetOptionalUint64()); - - std::cout << "Episode " << OptionalToString(parser.ColumnParser("episode_id").GetOptionalUint64()) - << ", Season: " << OptionalToString(parser.ColumnParser("season_id").GetOptionalUint64()) - << ", Title: " << OptionalToString(parser.ColumnParser("title").GetOptionalUtf8()) + TRowRange range(std::move(*resultSet)); + for (auto [episodeId, seasonId, title, airDateDays] : range.Get< + std::optional, + std::optional, + std::optional, + std::optional>({"episode_id", "season_id", "title", "air_date"})) { + auto airDate = TInstant::Days(*airDateDays); + + std::cout << "Episode " << OptionalToString(episodeId) + << ", Season: " << OptionalToString(seasonId) + << ", Title: " << OptionalToString(title) << ", Air date: " << airDate.FormatLocalTime("%a %b %d, %Y") << std::endl; } @@ -395,7 +399,7 @@ void ExplicitTcl(TQueryClient client) { void StreamQuerySelect(TQueryClient client) { std::cout << "> StreamQuery:" << std::endl; - ThrowOnError(client.RetryQuerySync([](TQueryClient client) -> TStatus { + ThrowOnError(client.RetryQuerySync([](TSession session) -> TStatus { auto query = R"( DECLARE $series AS List; @@ -419,44 +423,16 @@ void StreamQuerySelect(TQueryClient client) { .Build() .Build(); - // Executes stream query - auto resultStreamQuery = client.StreamExecuteQuery(query, TTxControl::NoTx(), parameters).GetValueSync(); - - if (!resultStreamQuery.IsSuccess()) { - return resultStreamQuery; - } - - // Iterates over results - bool eos = false; - - while (!eos) { - auto streamPart = resultStreamQuery.ReadNext().ExtractValueSync(); - - if (!streamPart.IsSuccess()) { - eos = true; - if (!streamPart.EOS()) { - return streamPart; - } - continue; - } - - // It is possible to duplicate lines in the output stream due to an external retryer. - if (streamPart.HasResultSet()) { - auto rs = streamPart.ExtractResultSet(); - TResultSetParser parser(rs); - while (parser.TryNextRow()) { - std::cout << "Season" - << ", SeriesId: " << OptionalToString(parser.ColumnParser("series_id").GetOptionalUint64()) - << ", SeasonId: " << OptionalToString(parser.ColumnParser("season_id").GetOptionalUint64()) - << ", Title: " << OptionalToString(parser.ColumnParser("title").GetOptionalUtf8()) - << ", Air date: " << parser.ColumnParser("first_aired").GetOptionalDate()->FormatLocalTime("%Y-%m-%d") - << std::endl; - } - } + for (auto& parser : TRowRange( + session.StreamExecuteQuery(query, TTxControl::NoTx(), parameters).ExtractValueSync())) { + std::cout << "Season" << ", SeriesId: " << OptionalToString(parser.ColumnParser("series_id").GetOptionalUint64()) + << ", SeasonId: " << OptionalToString(parser.ColumnParser("season_id").GetOptionalUint64()) + << ", Title: " << OptionalToString(parser.ColumnParser("title").GetOptionalUtf8()) + << ", Air date: " << parser.ColumnParser("first_aired").GetOptionalDate()->FormatLocalTime("%Y-%m-%d") + << std::endl; } return TStatus(EStatus::SUCCESS, NYdb::NIssue::TIssues()); })); - } diff --git a/include/ydb-cpp-sdk/client/iam/common/generic_provider.h b/include/ydb-cpp-sdk/client/iam/common/generic_provider.h index 5df5edd5eda..37d3184e4b6 100644 --- a/include/ydb-cpp-sdk/client/iam/common/generic_provider.h +++ b/include/ydb-cpp-sdk/client/iam/common/generic_provider.h @@ -168,7 +168,6 @@ class TGrpcIamCredentialsProvider : public ICredentialsProvider { } }; auto facility = weakFacility.lock(); - auto self = weakSelf.lock(); try { if (facility) { @@ -178,7 +177,7 @@ class TGrpcIamCredentialsProvider : public ICredentialsProvider { } catch (...) { } - if (self) { + if (auto self = weakSelf.lock()) { std::lock_guard guard(self->Lock_); self->ResetContextImpl(); } @@ -189,9 +188,13 @@ class TGrpcIamCredentialsProvider : public ICredentialsProvider { try { RequestFiller_(req); } catch (...) { + const auto now = SysClock::now(); std::lock_guard guard(Lock_); - LastRequestError_ = TStringBuilder() << "Request failed: " << CurrentExceptionMessage(); + LastRequestError_ = TStringBuilder() + << "Last request error was at " << FormatSysTimeUtcIsoMicros(now) + << ". Failed to prepare IAM request: " << CurrentExceptionMessage(); ResetContextImpl(); + RescheduleOnFailure(); return; } @@ -247,11 +250,23 @@ class TGrpcIamCredentialsProvider : public ICredentialsProvider { if (Context_.has_value() || SysClock::now() < NextTicketUpdate_) { return true; } - FillContext(guard); + try { + FillContext(guard); + } catch (...) { + const auto now = SysClock::now(); + LastRequestError_ = TStringBuilder() + << "Last request error was at " << FormatSysTimeUtcIsoMicros(now) + << ". Failed to prepare IAM request context: " << CurrentExceptionMessage(); + ResetContextImpl(); + } if (NeedStop_) { ResetContextImpl(); return false; } + if (!Context_.has_value()) { + RescheduleOnFailure(); + return true; + } } UpdateTicket(); return true; @@ -267,24 +282,13 @@ class TGrpcIamCredentialsProvider : public ICredentialsProvider { << " Message: \"" << status.error_message() << "\" iam-endpoint: \"" << IamEndpoint_.Endpoint << "\""; - const auto now = SysClock::now(); - const auto retryDelay = std::min(BackoffTimeout_, BACKOFF_MAX); - NextTicketUpdate_ = SafeAddSystemTime(now, ToBoundedSysDuration(retryDelay)); - BackoffTimeout_ = std::min(BackoffTimeout_ * 2, BACKOFF_MAX); + RescheduleOnFailure(); } else { LastRequestError_ = ""; Ticket_ = result.iam_token(); - BackoffTimeout_ = BACKOFF_START; - const auto now = SysClock::now(); - const SysTimePoint refreshAt = SafeAddSystemTime(now, ToBoundedSysDuration(IamEndpoint_.RefreshPeriod)); const SysTimePoint expiresAt = SysClock::from_time_t(result.expires_at().seconds()); - const SysDuration requestMargin = ToBoundedSysDuration(IamEndpoint_.RequestTimeout); - - SysTimePoint nextUpdate = std::min(refreshAt, expiresAt); - nextUpdate = SafeAddSystemTime(nextUpdate, -requestMargin); - nextUpdate = std::max(nextUpdate, SafeAddSystemTime(now, ToBoundedSysDuration(MINIMUM_REFRESH_INTERVAL))); - NextTicketUpdate_ = nextUpdate; + RescheduleOnSuccess(expiresAt); TokenReady_.notify_all(); } @@ -292,6 +296,26 @@ class TGrpcIamCredentialsProvider : public ICredentialsProvider { ResetContextImpl(); } + void RescheduleOnFailure() { // call with Lock_ + const auto now = SysClock::now(); + const auto retryDelay = std::min(BackoffTimeout_, BACKOFF_MAX); + NextTicketUpdate_ = SafeAddSystemTime(now, ToBoundedSysDuration(retryDelay)); + BackoffTimeout_ = std::min(BackoffTimeout_ * 2, BACKOFF_MAX); + } + + void RescheduleOnSuccess(const SysTimePoint expiresAt) { // call with Lock_ + BackoffTimeout_ = BACKOFF_START; + + const auto now = SysClock::now(); + const SysTimePoint refreshAt = SafeAddSystemTime(now, ToBoundedSysDuration(IamEndpoint_.RefreshPeriod)); + const SysDuration requestMargin = ToBoundedSysDuration(IamEndpoint_.RequestTimeout); + + SysTimePoint nextUpdate = std::min(refreshAt, expiresAt); + nextUpdate = SafeAddSystemTime(nextUpdate, -requestMargin); + nextUpdate = std::max(nextUpdate, SafeAddSystemTime(now, ToBoundedSysDuration(MINIMUM_REFRESH_INTERVAL))); + NextTicketUpdate_ = nextUpdate; + } + private: std::shared_ptr Channel_; std::shared_ptr Stub_; diff --git a/include/ydb-cpp-sdk/client/query/client.h b/include/ydb-cpp-sdk/client/query/client.h index ac6a6f41398..3b684dfb53d 100644 --- a/include/ydb-cpp-sdk/client/query/client.h +++ b/include/ydb-cpp-sdk/client/query/client.h @@ -62,12 +62,16 @@ struct TClientSettings : public TCommonClientSettingsBase { // ydb.query.session.pool.name. When empty the default // "@" is used. FLUENT_SETTING(std::string, PoolName); + + FLUENT_SETTING_DEFAULT(TRetryOperationSettings, RetrySettings, TRetryOperationSettings()); }; class TQueryClient { friend class TSession; friend class NRetry::Async::TRetryContext; friend class NRetry::Async::TRetryContext; + friend class NRetry::Async::TRetryContext>; + friend class NRetry::Async::TRetryContext; friend class NRetry::Sync::TRetryContext; public: @@ -110,13 +114,16 @@ class TQueryClient { TDuration timeout, bool isIndempotent); NThreading::TFuture ExecuteScript(const std::string& script, - const TExecuteScriptSettings& settings = TExecuteScriptSettings()); + const TExecuteScriptSettings& settings = TExecuteScriptSettings(), + const std::optional& retrySettings = std::nullopt); NThreading::TFuture ExecuteScript(const std::string& script, - const TParams& params, const TExecuteScriptSettings& settings = TExecuteScriptSettings()); + const TParams& params, const TExecuteScriptSettings& settings = TExecuteScriptSettings(), + const std::optional& retrySettings = std::nullopt); TAsyncFetchScriptResultsResult FetchScriptResults(const NKikimr::NOperationId::TOperationId& operationId, int64_t resultSetIndex, - const TFetchScriptResultsSettings& settings = TFetchScriptResultsSettings()); + const TFetchScriptResultsSettings& settings = TFetchScriptResultsSettings(), + const std::optional& retrySettings = std::nullopt); TAsyncCreateSessionResult GetSession(const TCreateSessionSettings& settings = TCreateSessionSettings()); @@ -129,6 +136,10 @@ class TQueryClient { //! Returns the size of session pool int64_t GetCurrentPoolSize() const; + // Internal: used by retry wrappers to suppress nested retries. + bool GetInRetryOperationContext() const; + void SetInRetryOperationContext(bool value); + private: class TImpl; std::shared_ptr Impl_; diff --git a/include/ydb-cpp-sdk/client/query/query.h b/include/ydb-cpp-sdk/client/query/query.h index ab127864955..23fc18c5558 100644 --- a/include/ydb-cpp-sdk/client/query/query.h +++ b/include/ydb-cpp-sdk/client/query/query.h @@ -5,6 +5,7 @@ #include "stats.h" #include +#include #include #include #include @@ -14,6 +15,8 @@ namespace NYdb::inline V3::NQuery { +using TRetryOperationSettings = NYdb::NRetry::TRetryOperationSettings; + enum class ESyntax { Unspecified = 0, YqlV1 = 1, // YQL @@ -113,6 +116,7 @@ struct TExecuteQuerySettings : public TRequestSettings { FLUENT_SETTING_DEFAULT(ESchemaInclusionMode, SchemaInclusionMode, ESchemaInclusionMode::Unspecified); FLUENT_SETTING_DEFAULT(TResultSet::EFormat, Format, TResultSet::EFormat::Unspecified); FLUENT_SETTING_OPTIONAL(TArrowFormatSettings, ArrowFormatSettings); + FLUENT_SETTING_OPTIONAL(TRetryOperationSettings, RetrySettings); }; struct TBeginTxSettings : public TRequestSettings {}; @@ -135,6 +139,7 @@ struct TExecuteScriptSettings : public TOperationRequestSettings { FLUENT_SETTING(std::string, FetchToken); FLUENT_SETTING_DEFAULT(uint64_t, RowsLimit, 1000); + FLUENT_SETTING_OPTIONAL(TRetryOperationSettings, RetrySettings); }; class TFetchScriptResultsResult : public TStatus { diff --git a/include/ydb-cpp-sdk/client/query/stats.h b/include/ydb-cpp-sdk/client/query/stats.h index 034fb0c5949..a1c944eb710 100644 --- a/include/ydb-cpp-sdk/client/query/stats.h +++ b/include/ydb-cpp-sdk/client/query/stats.h @@ -5,9 +5,14 @@ #include #include #include +#include namespace Ydb::TableStats { + class CompilationStats; + class OperationStats; + class QueryPhaseStats; class QueryStats; + class TableAccessStats; } namespace NYdb::inline V3 { @@ -16,6 +21,72 @@ namespace NYdb::inline V3 { namespace NYdb::inline V3::NQuery { +class TOperationStats { +public: + explicit TOperationStats(const Ydb::TableStats::OperationStats& proto); + + uint64_t GetRows() const; + uint64_t GetBytes() const; + +private: + uint64_t Rows_ = 0; + uint64_t Bytes_ = 0; +}; + +class TTableAccessStats { +public: + explicit TTableAccessStats(const Ydb::TableStats::TableAccessStats& proto); + + const std::string& GetName() const; + const TOperationStats& GetReads() const; + const TOperationStats& GetUpdates() const; + const TOperationStats& GetDeletes() const; + uint64_t GetPartitionsCount() const; + +private: + std::string Name_; + TOperationStats Reads_; + TOperationStats Updates_; + TOperationStats Deletes_; + uint64_t PartitionsCount_ = 0; +}; + +class TQueryPhaseStats { +public: + explicit TQueryPhaseStats(const Ydb::TableStats::QueryPhaseStats& proto); + + uint64_t GetDurationUs() const; + TDuration GetDuration() const; + uint64_t GetCpuTimeUs() const; + TDuration GetCpuTime() const; + uint64_t GetAffectedShards() const; + bool IsLiteralPhase() const; + const std::vector& GetTableAccess() const; + +private: + uint64_t DurationUs_ = 0; + uint64_t CpuTimeUs_ = 0; + uint64_t AffectedShards_ = 0; + bool LiteralPhase_ = false; + std::vector TableAccess_; +}; + +class TCompilationStats { +public: + explicit TCompilationStats(const Ydb::TableStats::CompilationStats& proto); + + bool IsFromCache() const; + uint64_t GetDurationUs() const; + TDuration GetDuration() const; + uint64_t GetCpuTimeUs() const; + TDuration GetCpuTime() const; + +private: + bool FromCache_ = false; + uint64_t DurationUs_ = 0; + uint64_t CpuTimeUs_ = 0; +}; + class TExecStats { friend class NYdb::TProtoAccessor; @@ -27,13 +98,20 @@ class TExecStats { std::string ToString(bool withPlan = false) const; + uint64_t GetProcessCpuTimeUs() const; + uint64_t GetTotalDurationUs() const; + uint64_t GetTotalCpuTimeUs() const; std::optional GetPlan() const; std::optional GetAst() const; std::optional GetMeta() const; + TDuration GetProcessCpuTime() const; TDuration GetTotalDuration() const; TDuration GetTotalCpuTime() const; + std::vector GetQueryPhases() const; + std::optional GetCompilation() const; + private: const Ydb::TableStats::QueryStats& GetProto() const; diff --git a/include/ydb-cpp-sdk/client/query/tx.h b/include/ydb-cpp-sdk/client/query/tx.h index 20167130f0d..843d4ecd0bc 100644 --- a/include/ydb-cpp-sdk/client/query/tx.h +++ b/include/ydb-cpp-sdk/client/query/tx.h @@ -44,6 +44,10 @@ struct TTxSettings { return TTxSettings(TS_READ_COMMITTED_RW); } + static TTxSettings StrictSerializableRW() { + return TTxSettings(TS_STRICT_SERIALIZABLE_RW); + } + void Out(IOutputStream& out) const { switch (Mode_) { case TS_SERIALIZABLE_RW: @@ -64,6 +68,9 @@ struct TTxSettings { case TS_READ_COMMITTED_RW: out << "ReadCommittedRW"; break; + case TS_STRICT_SERIALIZABLE_RW: + out << "StrictSerializableRW"; + break; default: out << "Unknown"; break; @@ -77,6 +84,7 @@ struct TTxSettings { TS_SNAPSHOT_RO, TS_SNAPSHOT_RW, TS_READ_COMMITTED_RW, + TS_STRICT_SERIALIZABLE_RW, }; FLUENT_SETTING(TTxOnlineSettings, OnlineSettings); diff --git a/include/ydb-cpp-sdk/client/result/detail/rows_column_getter.h b/include/ydb-cpp-sdk/client/result/detail/rows_column_getter.h new file mode 100644 index 00000000000..6332b1a9c86 --- /dev/null +++ b/include/ydb-cpp-sdk/client/result/detail/rows_column_getter.h @@ -0,0 +1,25 @@ +#pragma once + +// Internal implementation for TRowRange::Get. Not a supported extension point. + +namespace NYdb::inline V3 { + +namespace NRowRangesDetail { + +template +struct TRowColumnGetter { + static T Get(TRowParser& row, const std::string& name) { + return TValueParserGetter::Get(row.ColumnParser(name)); + } +}; + +template <> +struct TRowColumnGetter { + static TValue Get(TRowParser& row, const std::string& name) { + return row.GetValue(name); + } +}; + +} // namespace NRowRangesDetail + +} // namespace NYdb::inline V3 diff --git a/include/ydb-cpp-sdk/client/result/detail/rows_parser_getter.h b/include/ydb-cpp-sdk/client/result/detail/rows_parser_getter.h new file mode 100644 index 00000000000..7dff811550e --- /dev/null +++ b/include/ydb-cpp-sdk/client/result/detail/rows_parser_getter.h @@ -0,0 +1,378 @@ +#pragma once + +// Internal implementation for TRowRange::Get. Not a supported extension point. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace NYdb::inline V3 { + +namespace NRowRangesDetail { + +namespace NParserScope { + +template +struct TParserScope { + explicit TParserScope(TValueParser& parser) + : Parser_(parser) + { + (Parser_.*Open)(); + } + + ~TParserScope() { + (Parser_.*Close)(); + } + + TParserScope(const TParserScope&) = delete; + TParserScope& operator=(const TParserScope&) = delete; + + TValueParser& Parser_; +}; + +using TOptionalScope = TParserScope<&TValueParser::OpenOptional, &TValueParser::CloseOptional>; +using TListScope = TParserScope<&TValueParser::OpenList, &TValueParser::CloseList>; +using TDictScope = TParserScope<&TValueParser::OpenDict, &TValueParser::CloseDict>; +using TTupleScope = TParserScope<&TValueParser::OpenTuple, &TValueParser::CloseTuple>; +using TStructScope = TParserScope<&TValueParser::OpenStruct, &TValueParser::CloseStruct>; +using TVariantScope = TParserScope<&TValueParser::OpenVariant, &TValueParser::CloseVariant>; +using TTaggedScope = TParserScope<&TValueParser::OpenTagged, &TValueParser::CloseTagged>; + +} // namespace NParserScope + +//! Maps a C++ type to the TValueParser accessor that reads it. +//! Unspecialised T fails to compile ("incomplete type") — that is the intended +//! error for unsupported column types. +template +struct TValueParserGetter; + +using TDate32 = std::chrono::sys_time; +using TDatetime64 = std::chrono::sys_time; +using TTimestamp64 = std::chrono::sys_time; + +//! Any default-constructible sequence container with push_back/emplace_back. +template +concept CRowSequenceContainer = requires(T container, typename T::value_type value) { + typename T::value_type; + { T{} } -> std::same_as; + requires requires { container.push_back(value); } || + requires { container.emplace_back(value); }; +}; + +//! Any default-constructible associative container with emplace or insert. +template +concept CRowAssociativeContainer = requires(T container, typename T::key_type key, typename T::mapped_type mapped) { + typename T::key_type; + typename T::mapped_type; + { T{} } -> std::same_as; + requires requires { container.emplace(key, mapped); } || + requires { container.insert({key, mapped}); }; +}; + +template +struct TValueParserGetter> { + static std::optional Get(TValueParser& p) { + NParserScope::TOptionalScope scope(p); + if (p.IsNull()) { + return std::nullopt; + } + return TValueParserGetter::Get(p); + } +}; + +#define Y_DEFINE_VALUE_PARSER_GETTER(Type, Suffix) \ + template <> struct TValueParserGetter { \ + static Type Get(TValueParser& p) { return p.Get##Suffix(); } \ + } + +Y_DEFINE_VALUE_PARSER_GETTER(bool, Bool); +Y_DEFINE_VALUE_PARSER_GETTER(int8_t, Int8); +Y_DEFINE_VALUE_PARSER_GETTER(uint8_t, Uint8); +Y_DEFINE_VALUE_PARSER_GETTER(int16_t, Int16); +Y_DEFINE_VALUE_PARSER_GETTER(uint16_t, Uint16); +Y_DEFINE_VALUE_PARSER_GETTER(int32_t, Int32); +Y_DEFINE_VALUE_PARSER_GETTER(uint32_t, Uint32); +Y_DEFINE_VALUE_PARSER_GETTER(uint64_t, Uint64); +Y_DEFINE_VALUE_PARSER_GETTER(float, Float); +Y_DEFINE_VALUE_PARSER_GETTER(double, Double); +Y_DEFINE_VALUE_PARSER_GETTER(TUuidValue, Uuid); +Y_DEFINE_VALUE_PARSER_GETTER(TDate32, Date32); +Y_DEFINE_VALUE_PARSER_GETTER(TDatetime64, Datetime64); +Y_DEFINE_VALUE_PARSER_GETTER(TTimestamp64, Timestamp64); +Y_DEFINE_VALUE_PARSER_GETTER(TWideMicroseconds, Interval64); +Y_DEFINE_VALUE_PARSER_GETTER(TDecimalValue, Decimal); +Y_DEFINE_VALUE_PARSER_GETTER(TPgValue, Pg); + +#undef Y_DEFINE_VALUE_PARSER_GETTER + +template <> struct TValueParserGetter { + static int64_t Get(TValueParser& p) { + switch (p.GetPrimitiveType()) { + case EPrimitiveType::Int64: return p.GetInt64(); + case EPrimitiveType::Interval: return p.GetInterval(); + default: + throw std::runtime_error( + "TValueParserGetter: column type is not Int64/Interval"); + } + } +}; + +template <> struct TValueParserGetter { + static std::string Get(TValueParser& p) { + switch (p.GetPrimitiveType()) { + case EPrimitiveType::String: return p.GetString(); + case EPrimitiveType::Utf8: return p.GetUtf8(); + case EPrimitiveType::Yson: return p.GetYson(); + case EPrimitiveType::Json: return p.GetJson(); + case EPrimitiveType::JsonDocument: return p.GetJsonDocument(); + case EPrimitiveType::DyNumber: return p.GetDyNumber(); + case EPrimitiveType::TzDate: return p.GetTzDate(); + case EPrimitiveType::TzDatetime: return p.GetTzDatetime(); + case EPrimitiveType::TzTimestamp: return p.GetTzTimestamp(); + default: + throw std::runtime_error( + "TValueParserGetter: column type is not a string-like primitive"); + } + } +}; + +//! TInstant maps to whichever date/time primitive the column actually holds: +//! Date, Datetime, or Timestamp. +template <> struct TValueParserGetter { + static TInstant Get(TValueParser& p) { + switch (p.GetPrimitiveType()) { + case EPrimitiveType::Date: return p.GetDate(); + case EPrimitiveType::Datetime: return p.GetDatetime(); + case EPrimitiveType::Timestamp: return p.GetTimestamp(); + default: + throw std::runtime_error( + "TValueParserGetter: column type is not Date/Datetime/Timestamp"); + } + } +}; + +namespace NContainerDetail { + +template +void PushBack(Sequence& sequence, Value&& value) { + if constexpr (requires { + sequence.emplace_back(std::forward(value)); + }) { + sequence.emplace_back(std::forward(value)); + } else { + sequence.push_back(std::forward(value)); + } +} + +template +void EmplaceMapEntry(Map& map, Key&& key, Value&& value) { + if constexpr (requires { + map.emplace(std::forward(key), std::forward(value)); + }) { + map.emplace(std::forward(key), std::forward(value)); + } else { + map.insert({std::forward(key), std::forward(value)}); + } +} + +template +Sequence GetList(TValueParser& p) { + using Element = typename Sequence::value_type; + Sequence result; + NParserScope::TListScope scope(p); + while (p.TryNextListItem()) { + PushBack(result, TValueParserGetter::Get(p)); + } + return result; +} + +template +Map GetDict(TValueParser& p) { + using Key = typename Map::key_type; + using Value = typename Map::mapped_type; + + Map result; + NParserScope::TDictScope scope(p); + while (p.TryNextDictItem()) { + p.DictKey(); + Key key = TValueParserGetter::Get(p); + p.DictPayload(); + Value value = TValueParserGetter::Get(p); + EmplaceMapEntry(result, std::move(key), std::move(value)); + } + return result; +} + +} // namespace NContainerDetail + +template +struct TValueParserGetter { + static Sequence Get(TValueParser& p) { + if (p.GetKind() != TTypeParser::ETypeKind::List) { + throw std::runtime_error("TValueParserGetter: column type is not List"); + } + return NContainerDetail::GetList(p); + } +}; + +template +struct TValueParserGetter { + static Map Get(TValueParser& p) { + if (p.GetKind() != TTypeParser::ETypeKind::Dict) { + throw std::runtime_error("TValueParserGetter: column type is not Dict"); + } + return NContainerDetail::GetDict(p); + } +}; + +namespace NTupleDetail { + +template +void GetTupleElement(TValueParser& p, Tuple& result, bool (*TryNext)(TValueParser&)) { + if (!TryNext(p)) { + throw std::runtime_error("TValueParserGetter>: not enough elements"); + } + std::get(result) = TValueParserGetter>::Get(p); +} + +template +Tuple GetTupleElements(TValueParser& p, std::index_sequence, bool (*TryNext)(TValueParser&)) { + Tuple result; + (GetTupleElement(p, result, TryNext), ...); + return result; +} + +inline void CheckNoExtraElements(TValueParser& p, bool (*TryNext)(TValueParser&)) { + if (TryNext(p)) { + throw std::runtime_error("TValueParserGetter>: too many elements"); + } +} + +inline bool TryNextTupleElement(TValueParser& p) { + return p.TryNextElement(); +} + +inline bool TryNextStructMember(TValueParser& p) { + return p.TryNextMember(); +} + +template +std::tuple GetTupleFromContainer( + TValueParser& p, + TTypeParser::ETypeKind expectedKind, + bool (*TryNext)(TValueParser&)) { + if (p.GetKind() != expectedKind) { + throw std::runtime_error("TValueParserGetter>: column type mismatch"); + } + if (expectedKind == TTypeParser::ETypeKind::Tuple) { + NParserScope::TTupleScope scope(p); + auto result = GetTupleElements>( + p, std::index_sequence_for{}, TryNext); + CheckNoExtraElements(p, TryNext); + return result; + } + NParserScope::TStructScope scope(p); + auto result = GetTupleElements>( + p, std::index_sequence_for{}, TryNext); + CheckNoExtraElements(p, TryNext); + return result; +} + +} // namespace NTupleDetail + +template +struct TValueParserGetter> { + static std::tuple Get(TValueParser& p) { + switch (p.GetKind()) { + case TTypeParser::ETypeKind::Tuple: + return NTupleDetail::GetTupleFromContainer( + p, TTypeParser::ETypeKind::Tuple, NTupleDetail::TryNextTupleElement); + case TTypeParser::ETypeKind::Struct: + return NTupleDetail::GetTupleFromContainer( + p, TTypeParser::ETypeKind::Struct, NTupleDetail::TryNextStructMember); + default: + throw std::runtime_error( + "TValueParserGetter>: column type is not Tuple/Struct"); + } + } +}; + +namespace NVariantDetail { + +template +bool TryEmplaceVariantAlternative(Variant& result, TValueParser& p, size_t index) { + if (index != I) { + return false; + } + result.template emplace( + TValueParserGetter>::Get(p)); + return true; +} + +template +bool EmplaceVariantAlternative(Variant& result, TValueParser& p, size_t index, std::index_sequence) { + return (TryEmplaceVariantAlternative(result, p, index) || ...); +} + +} // namespace NVariantDetail + +template +struct TValueParserGetter> { + static std::variant Get(TValueParser& p) { + if (p.GetKind() != TTypeParser::ETypeKind::Variant) { + throw std::runtime_error("TValueParserGetter>: column type is not Variant"); + } + const size_t index = p.GetVariantIndex(); + NParserScope::TVariantScope scope(p); + std::variant result; + if (!NVariantDetail::EmplaceVariantAlternative( + result, p, index, std::index_sequence_for{})) { + throw std::runtime_error("TValueParserGetter>: variant index out of range"); + } + return result; + } +}; + +template +struct TValueParserGetter> { + static std::pair Get(TValueParser& p) { + if (p.GetKind() != TTypeParser::ETypeKind::Tagged) { + throw std::runtime_error( + "TValueParserGetter>: column type is not Tagged"); + } + NParserScope::TTaggedScope scope(p); + std::string tag = p.GetTag(); + T value = TValueParserGetter::Get(p); + return {std::move(tag), std::move(value)}; + } +}; + +template +std::array BuildColumnNames(It first, It last) { + std::array names; + size_t i = 0; + for (auto it = first; it != last; ++it) { + if (i >= N) { + throw std::invalid_argument( + "TRowRange::Get: too many column names for the requested types"); + } + names[i++] = std::string(*it); + } + if (i != N) { + throw std::invalid_argument( + "TRowRange::Get: not enough column names for the requested types"); + } + return names; +} + +} // namespace NRowRangesDetail +} // namespace NYdb::inline V3 diff --git a/include/ydb-cpp-sdk/client/result/rows.h b/include/ydb-cpp-sdk/client/result/rows.h new file mode 100644 index 00000000000..96d42416dc1 --- /dev/null +++ b/include/ydb-cpp-sdk/client/result/rows.h @@ -0,0 +1,245 @@ +#pragma once + +#include "result.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "detail/rows_parser_getter.h" + +namespace NYdb::inline V3 { + +class TRowIterEnd {}; + +class TRowParser; + +class TRowParserHolder; + +template +class TRowColumns; + +class TRowIterator; + +class TRowRange; + +class TRowParser { + friend class TRowParserHolder; +public: + TRowParser(const TRowParser&) = delete; + TRowParser(TRowParser&&) = delete; + TRowParser& operator=(const TRowParser&) = delete; + TRowParser& operator=(TRowParser&&) = delete; + + size_t ColumnsCount() const; + size_t RowsCount() const; + ssize_t ColumnIndex(const std::string& columnName); + TValueParser& ColumnParser(size_t columnIndex); + TValueParser& ColumnParser(const std::string& columnName); + TValue GetValue(size_t columnIndex) const; + TValue GetValue(const std::string& columnName) const; + +private: + explicit TRowParser(TResultSetParser& parser); + + TResultSetParser& Parser_; +}; + +} // namespace NYdb::inline V3 + +#include "detail/rows_column_getter.h" + +namespace NYdb::inline V3 { + +//! Move-only, single-pass input iterator over rows. +//! Valid while shared range state is alive. Not copyable and not multipass: +//! do not expect independent copies, re-traversal, or classic STL algorithms +//! that require copyable input iterators. Advertises input_iterator_tag for +//! ranges-style sentinel iteration, not classic copyable input iterators. +class TRowIterator { + friend class TRowRange; +public: + class TImpl; + + using iterator_category = std::input_iterator_tag; + using value_type = TRowParser; + using difference_type = std::ptrdiff_t; + using pointer = TRowParser*; + using reference = TRowParser&; + + TRowIterator(const TRowIterator&) = delete; + TRowIterator(TRowIterator&&) noexcept; + ~TRowIterator(); + TRowIterator& operator=(const TRowIterator&) = delete; + TRowIterator& operator=(TRowIterator&&) noexcept; + + bool operator==(const TRowIterEnd& end) const noexcept; + bool operator!=(const TRowIterEnd& end) const noexcept; + + reference operator*() const; + pointer operator->() const; + + TRowIterator& operator++(); + void operator++(int) { + ++(*this); + } + +private: + explicit TRowIterator(std::unique_ptr impl); + + std::unique_ptr Impl_; +}; + +inline bool operator==(const TRowIterEnd& end, const TRowIterator& it) noexcept { + return it == end; +} + +inline bool operator!=(const TRowIterEnd& end, const TRowIterator& it) noexcept { + return it != end; +} + +//! Forward, single-pass range over rows of a single logical result set. +class TRowRange { + friend class TRowIterator::TImpl; + template + friend class TRowColumns; +public: + explicit TRowRange(TResultSet&& resultSet); + explicit TRowRange(NTable::TDataQueryResult&& result); + + explicit TRowRange(NQuery::TExecuteQueryIterator&& iterator); + explicit TRowRange(NTable::TScanQueryPartIterator&& iterator); + explicit TRowRange(NTable::TTablePartIterator&& iterator); + + TRowRange(const TRowRange&) = delete; + TRowRange(TRowRange&&) noexcept; + TRowRange& operator=(const TRowRange&) = delete; + TRowRange& operator=(TRowRange&&) noexcept; + ~TRowRange(); + + TRowIterator begin(); + TRowIterEnd end() const noexcept; + + //! Build a typed-tuple view over this range. + //! Args... are the column C++ types (e.g. int32_t, std::optional); + //! columns is any iterable of names convertible to std::string. + //! Yields std::tuple per row. + template + TRowColumns Get(std::initializer_list columns); + + template + TRowColumns Get(const StringIterable& columns); + + class TImpl; + +private: + static TRowIterator BeginIterator(const std::shared_ptr& state); + + std::shared_ptr Impl_; +}; + +//! Lightweight typed view over a TRowRange. Holds shared range state and column +//! names; yields std::tuple per row. +template +class TRowColumns { + friend class TRowRange; + static constexpr size_t N = sizeof...(Args); + using TNames = std::array; +public: + class Iterator { + friend class TRowColumns; + public: + using iterator_category = std::input_iterator_tag; + using value_type = std::tuple; + using difference_type = std::ptrdiff_t; + using pointer = void; + using reference = value_type; + + Iterator(const Iterator&) = delete; + Iterator(Iterator&&) noexcept = default; + Iterator& operator=(const Iterator&) = delete; + Iterator& operator=(Iterator&&) noexcept = default; + + value_type operator*() const { + return MakeTuple(std::index_sequence_for{}); + } + + Iterator& operator++() { + ++Inner_; + return *this; + } + + void operator++(int) { + ++(*this); + } + + bool operator==(const TRowIterEnd& end) const { return Inner_ == end; } + bool operator!=(const TRowIterEnd& end) const { return Inner_ != end; } + + private: + Iterator(TRowIterator&& inner, const TNames& names) + : Inner_(std::move(inner)) + , Names_(&names) + {} + + template + value_type MakeTuple(std::index_sequence) const { + TRowParser& row = *Inner_; + return value_type{ + NRowRangesDetail::TRowColumnGetter::Get(row, (*Names_)[Is])... + }; + } + + TRowIterator Inner_; + const TNames* Names_; + }; + + TRowColumns(const TRowColumns&) = delete; + TRowColumns(TRowColumns&&) noexcept = default; + TRowColumns& operator=(const TRowColumns&) = delete; + TRowColumns& operator=(TRowColumns&&) noexcept = default; + + Iterator begin() { + return Iterator(TRowRange::BeginIterator(State_), Names_); + } + + TRowIterEnd end() const noexcept { + return TRowIterEnd{}; + } + +private: + TRowColumns(std::shared_ptr state, TNames&& names) + : State_(std::move(state)) + , Names_(std::move(names)) + {} + + std::shared_ptr State_; + TNames Names_; +}; + +template +TRowColumns TRowRange::Get(std::initializer_list columns) { + return TRowColumns( + Impl_, + NRowRangesDetail::BuildColumnNames(columns.begin(), columns.end())); +} + +template +TRowColumns TRowRange::Get(const StringIterable& columns) { + return TRowColumns( + Impl_, + NRowRangesDetail::BuildColumnNames(std::begin(columns), std::end(columns))); +} + +} // namespace NYdb::inline V3 diff --git a/include/ydb-cpp-sdk/client/retry/retry.h b/include/ydb-cpp-sdk/client/retry/retry.h index 3c656a6a84c..62003137ac1 100644 --- a/include/ydb-cpp-sdk/client/retry/retry.h +++ b/include/ydb-cpp-sdk/client/retry/retry.h @@ -22,7 +22,13 @@ struct TRetryOperationSettings { FLUENT_SETTING_DEFAULT(TDuration, MaxTimeout, TDuration::Max()); FLUENT_SETTING_DEFAULT(TBackoffSettings, FastBackoffSettings, DefaultFastBackoffSettings()); FLUENT_SETTING_DEFAULT(TBackoffSettings, SlowBackoffSettings, DefaultSlowBackoffSettings()); - FLUENT_SETTING_FLAG(Idempotent); + bool Idempotent_ = false; + bool IdempotentWasSet_ = false; + TSelf& Idempotent(bool value = true) { + Idempotent_ = value; + IdempotentWasSet_ = true; + return static_cast(*this); + } FLUENT_SETTING_FLAG(Verbose); FLUENT_SETTING_FLAG(RetryUndefined); diff --git a/include/ydb-cpp-sdk/client/table/fwd.h b/include/ydb-cpp-sdk/client/table/fwd.h index f10a5f29048..8315bce41e4 100644 --- a/include/ydb-cpp-sdk/client/table/fwd.h +++ b/include/ydb-cpp-sdk/client/table/fwd.h @@ -65,6 +65,7 @@ struct TReplicationPolicy; class TBuildIndexOperation; class TCompactionOperation; +class TAnalyzeOperation; class TTtlDeleteAction; class TTtlEvictToExternalStorageAction; diff --git a/include/ydb-cpp-sdk/client/table/table.h b/include/ydb-cpp-sdk/client/table/table.h index 9c171202a0c..fa892c3169f 100644 --- a/include/ydb-cpp-sdk/client/table/table.h +++ b/include/ydb-cpp-sdk/client/table/table.h @@ -13,6 +13,7 @@ #include #include +#include #include namespace Ydb { @@ -64,6 +65,7 @@ class TRetryContext; namespace NRetry { template class TRetryDeadlineHelper; +class TBulkUpsertRetryState; } // namespace NRetry namespace NScheme { @@ -603,6 +605,24 @@ class TCompactionOperation : public TOperation { TMetadata Metadata_; }; +class TAnalyzeOperation : public TOperation { +public: + using TOperation::TOperation; + TAnalyzeOperation(TStatus&& status, Ydb::Operations::Operation&& operation); + + struct TMetadata { + EAnalyzeState State = EAnalyzeState::Unspecified; + float Progress = 0; + std::vector Paths; // All paths covered by this analyze. + std::vector InProgressPaths; // Subset of Paths currently being traversed + std::vector DonePaths; // Subset of Paths whose analysis completed + }; + + const TMetadata& Metadata() const; +private: + TMetadata Metadata_; +}; + //////////////////////////////////////////////////////////////////////////////// //! Represents changefeed description @@ -1545,17 +1565,24 @@ struct TClientSettings : public TCommonClientSettingsBase { // ydb.table.session.pool.name. When empty the default // "@" is used. FLUENT_SETTING(std::string, PoolName); + + FLUENT_SETTING_DEFAULT(TRetryOperationSettings, RetrySettings, TRetryOperationSettings()); }; struct TBulkUpsertSettings : public TOperationRequestSettings { // Format setting proto serialized into string. If not set format defaults are used. // I.e. it's Ydb.Table.CsvSettings for CSV. FLUENT_SETTING_DEFAULT(std::string, FormatSettings, ""); + FLUENT_SETTING_OPTIONAL(TRetryOperationSettings, RetrySettings); google::protobuf::Arena* Arena_ = nullptr; TBulkUpsertSettings& Arena(google::protobuf::Arena* arena) { Arena_ = arena; return *this; } + + // Internal: BulkUpsert retry backup hook. + std::shared_ptr RetryRowsState_; }; struct TReadRowsSettings : public TOperationRequestSettings { + FLUENT_SETTING_OPTIONAL(TRetryOperationSettings, RetrySettings); }; struct TStreamExecScanQuerySettings : public TRequestSettings { @@ -1581,6 +1608,8 @@ class TTableClient { friend class TSessionPool; friend class NRetry::Sync::TRetryContext; friend class NRetry::Async::TRetryContext; + friend class NRetry::Async::TRetryContext; + friend class NRetry::Async::TRetryContext; public: using TOperationFunc = std::function; @@ -1611,6 +1640,10 @@ class TTableClient { //! Returns the size of session pool int64_t GetCurrentPoolSize() const; + // Internal: used by retry wrappers to suppress nested retries. + bool GetInRetryOperationContext() const; + void SetInRetryOperationContext(bool value); + //! Returns new table builder TTableBuilder GetTableBuilder(); //! Returns new params builder @@ -1685,6 +1718,7 @@ struct TTxOnlineSettings { FLUENT_SETTING_DEFAULT(bool, AllowInconsistentReads, false); }; + class TTxSettings { friend class TTableClient; @@ -1714,6 +1748,10 @@ class TTxSettings { return TTxSettings(TS_SNAPSHOT_RW); } + static TTxSettings StrictSerializableRW() { + return TTxSettings(TS_STRICT_SERIALIZABLE_RW); + } + void Out(IOutputStream& out) const { switch (Mode_) { case TS_SERIALIZABLE_RW: @@ -1731,6 +1769,9 @@ class TTxSettings { case TS_SNAPSHOT_RW: out << "SnapshotRW"; break; + case TS_STRICT_SERIALIZABLE_RW: + out << "StrictSerializableRW"; + break; default: out << "Unknown"; break; @@ -1744,6 +1785,7 @@ class TTxSettings { TS_STALE_RO, TS_SNAPSHOT_RO, TS_SNAPSHOT_RW, + TS_STRICT_SERIALIZABLE_RW, }; FLUENT_SETTING(TTxOnlineSettings, OnlineSettings); @@ -2724,6 +2766,7 @@ class TReadRowsResult : public TStatus { TResultSet ResultSet; public: + explicit TReadRowsResult(TStatus&& status); explicit TReadRowsResult(TStatus&& status, TResultSet&& resultSet); TResultSet GetResultSet() { diff --git a/include/ydb-cpp-sdk/client/table/table_enum.h b/include/ydb-cpp-sdk/client/table/table_enum.h index dd10eaa274b..6ae6d50e41f 100644 --- a/include/ydb-cpp-sdk/client/table/table_enum.h +++ b/include/ydb-cpp-sdk/client/table/table_enum.h @@ -52,6 +52,15 @@ enum class ECompactState { Cancelled = 3, }; +enum class EAnalyzeState { + Unspecified = 0, + Enqueued = 1, + InProgress = 2, + Done = 3, + Cancelled = 4, + Failed = 5, +}; + enum class EChangefeedMode { KeysOnly /* "KEYS_ONLY" */, Updates /* "UPDATES" */, diff --git a/include/ydb-cpp-sdk/client/topic/codecs.h b/include/ydb-cpp-sdk/client/topic/codecs.h index 857ab0bd3ff..00c4e8a77f3 100644 --- a/include/ydb-cpp-sdk/client/topic/codecs.h +++ b/include/ydb-cpp-sdk/client/topic/codecs.h @@ -7,6 +7,12 @@ #include #include +#include +#include + +#include +#include +#include #include #include @@ -18,6 +24,7 @@ enum class ECodec : uint32_t { GZIP = 2, LZOP = 3, ZSTD = 4, + KAFKA_BATCH = 5, CUSTOM = 10000, }; @@ -32,11 +39,46 @@ inline const std::string& GetCodecId(const ECodec codec) { return idByCodec[codec]; } +struct TWriteBlockCompression { + ECodec Codec = ECodec::RAW; + std::vector& Payloads; + const std::vector& CreatedAt; + TBuffer& Data; + ui32& CodecID; + bool& Compressed; + i64 BaseSequence = 0; + std::optional BatchInnerCodec; + int CompressionLevel = 0; +}; + +//! Per-record metadata extracted from codec payload (e.g. Kafka record batch). +//! When not set, read path uses metadata from outer message_data. +struct TDecompressedMessageMeta { + i32 OffsetDelta; + i64 SequenceDelta; + i64 TimestampDelta; +}; + +struct TDecompressedMessage { + std::string Data; + std::optional Meta; +}; + +struct TDecompressionResult { + std::vector Messages; + //! Populated when codec expands one blob into many messages (Kafka batch). + std::optional BatchBaseOffset; + std::optional BatchBaseSequence; + std::optional BatchBaseTimestampMs; +}; + class ICodec { public: virtual ~ICodec() = default; virtual std::string Decompress(const std::string& data) const = 0; + virtual TDecompressionResult DecompressData(const std::string& data) const; virtual std::unique_ptr CreateCoder(TBuffer& result, int quality) const = 0; + virtual void CompressWriteBlock(TWriteBlockCompression& ctx) const; }; class TGzipCodec final : public ICodec { @@ -57,6 +99,17 @@ class TUnsupportedCodec final : public ICodec { std::unique_ptr CreateCoder(TBuffer&, int) const override; }; +class TKafkaBatchCodec final : public ICodec { +public: + std::string Decompress(const std::string& data) const override; + + TDecompressionResult DecompressData(const std::string& data) const override; + + std::unique_ptr CreateCoder(TBuffer&, int) const override; + + void CompressWriteBlock(TWriteBlockCompression& ctx) const override; +}; + class TCodecMap { public: static TCodecMap& GetTheCodecMap() { @@ -93,4 +146,11 @@ class TCodecMap { TAdaptiveLock Lock; }; +inline std::string TakeFirstDecompressedMessage(TDecompressionResult&& result) { + if (result.Messages.empty()) { + throw yexception() << "empty decompression result"; + } + return std::move(result.Messages.front().Data); +} + } // namespace NYdb::NTopic diff --git a/include/ydb-cpp-sdk/client/topic/read_events.h b/include/ydb-cpp-sdk/client/topic/read_events.h index 081b2b5850a..d9b2369187a 100644 --- a/include/ydb-cpp-sdk/client/topic/read_events.h +++ b/include/ydb-cpp-sdk/client/topic/read_events.h @@ -63,7 +63,7 @@ struct TPartitionSessionControl: public TPartitionSession { virtual void Commit(uint64_t startOffset, uint64_t endOffset) = 0; //! Confirm partition session creation from TStartPartitionSessionEvent. - virtual void ConfirmCreate(std::optional readOffset, std::optional commitOffset) = 0; + virtual void ConfirmCreate(std::optional readOffset, std::optional commitOffset, std::optional maxOffset = std::nullopt) = 0; //! Confirm partition session destruction from TStopPartitionSessionEvent. virtual void ConfirmDestroy() = 0; @@ -98,7 +98,8 @@ struct TReadSessionEvent { TWriteSessionMeta::TPtr meta, TMessageMeta::TPtr messageMeta, uint64_t uncompressedSize, - std::string messageGroupId); + std::string messageGroupId, + uint64_t logicalMessageCount = 1); uint64_t Offset; std::string ProducerId; uint64_t SeqNo; @@ -108,6 +109,7 @@ struct TReadSessionEvent { TMessageMeta::TPtr MessageMeta; uint64_t UncompressedSize; std::string MessageGroupId; + uint64_t LogicalMessageCount; }; class TMessageBase: public TPrintable { @@ -124,6 +126,9 @@ struct TReadSessionEvent { //! Message offset. uint64_t GetOffset() const; + //! Number of logical messages covered by this message. + uint64_t GetLogicalMessageCount() const; + //! Producer id. const std::string& GetProducerId() const; @@ -291,7 +296,7 @@ struct TReadSessionEvent { //! Confirm partition session creation. //! This signals that user is ready to receive data from this partition session. //! If maybe is empty then no rewinding - void Confirm(std::optional readOffset = std::nullopt, std::optional commitOffset = std::nullopt); + void Confirm(std::optional readOffset = std::nullopt, std::optional commitOffset = std::nullopt, std::optional maxOffset = std::nullopt); private: uint64_t CommittedOffset; diff --git a/include/ydb-cpp-sdk/client/topic/write_session.h b/include/ydb-cpp-sdk/client/topic/write_session.h index d25e9e80e7b..0eb05c9dffc 100644 --- a/include/ydb-cpp-sdk/client/topic/write_session.h +++ b/include/ydb-cpp-sdk/client/topic/write_session.h @@ -98,6 +98,10 @@ struct TWriteSessionSettings : public TRequestSettings { FLUENT_SETTING_DEFAULT(ECodec, Codec, ECodec::GZIP); FLUENT_SETTING_DEFAULT(int32_t, CompressionLevel, 4); + //! Inner compression for Kafka record batch payload when Codec is KAFKA_BATCH. + //! Allowed values: not set (no inner compression), GZIP, ZSTD. + FLUENT_SETTING_OPTIONAL(ECodec, BatchInnerCodec); + //! Writer will not accept new messages if memory usage exceeds this limit. //! Memory usage consists of raw data pending compression and compressed messages being sent. FLUENT_SETTING_DEFAULT(uint64_t, MaxMemoryUsage, 20_MB); @@ -122,11 +126,14 @@ struct TWriteSessionSettings : public TRequestSettings { //! but for no longer than BatchFlushInterval. //! Upon reaching FlushInterval or FlushSize limit, all messages will be written with one batch. //! Greatly increases performance for small messages. - //! Setting either value to zero means immediate write with no batching. (Unrecommended, especially for clients - //! sending small messages at high rate). - FLUENT_SETTING_OPTIONAL(TDuration, BatchFlushInterval); + //! Zero BatchFlushInterval or BatchFlushSizeBytes disables the corresponding limit (immediate flush). + FLUENT_SETTING_DEFAULT(TDuration, BatchFlushInterval, TDuration::Seconds(1)); FLUENT_SETTING_OPTIONAL(uint64_t, BatchFlushSizeBytes); + //! Max number of logical messages packed into a single write block. + //! Values greater than 1 are sent as a single batch block. + FLUENT_SETTING_DEFAULT(uint32_t, BatchFlushMessageCount, 1); + FLUENT_SETTING_DEFAULT(TDuration, ConnectTimeout, TDuration::Seconds(30)); FLUENT_SETTING_OPTIONAL(TWriterCounters::TPtr, Counters); diff --git a/include/ydb-cpp-sdk/client/types/fwd.h b/include/ydb-cpp-sdk/client/types/fwd.h index 6168920659e..63b1c89aaeb 100644 --- a/include/ydb-cpp-sdk/client/types/fwd.h +++ b/include/ydb-cpp-sdk/client/types/fwd.h @@ -29,6 +29,7 @@ class ITokenSource; namespace NStatusHelpers { class TYdbErrorException; +class TYdbRangeErrorException; } } // namespace NYdb diff --git a/include/ydb-cpp-sdk/client/types/status/status.h b/include/ydb-cpp-sdk/client/types/status/status.h index 0e86eea0aa0..cbd1f53a141 100644 --- a/include/ydb-cpp-sdk/client/types/status/status.h +++ b/include/ydb-cpp-sdk/client/types/status/status.h @@ -73,10 +73,34 @@ class TYdbErrorException : public TYdbException { TStatus Status_; }; +class TYdbRangeErrorException : public TYdbException { +public: + TYdbRangeErrorException(TStatus status) + : Status_(std::move(status)) + { + } + + friend IOutputStream& operator<<(IOutputStream& out, const TYdbRangeErrorException& e) { + return out << e.Status_; + } + + const TStatus& GetStatus() const { + return Status_; + } + + TStatus&& ExtractStatus() { + return std::move(Status_); + } + +private: + TStatus Status_; +}; + void ThrowOnError(TStatus status, std::function onSuccess = [](TStatus) {}); void ThrowOnErrorOrPrintIssues(TStatus status); bool StatusContainsIssueWithCode(const TStatus& status, NYdb::NIssue::TIssueCode code); + } } // namespace NYdb diff --git a/include/ydb-cpp-sdk/client/value/value.h b/include/ydb-cpp-sdk/client/value/value.h index 66cba7a6b81..4fafc0710eb 100644 --- a/include/ydb-cpp-sdk/client/value/value.h +++ b/include/ydb-cpp-sdk/client/value/value.h @@ -413,6 +413,7 @@ class TValueParser : public TMoveOnly { void CloseDict(); // Variant + uint32_t GetVariantIndex() const; void OpenVariant(); void CloseVariant(); diff --git a/include/ydb-cpp-sdk/library/operation_id/operation_id.h b/include/ydb-cpp-sdk/library/operation_id/operation_id.h index 8677e198680..57b69eb2672 100644 --- a/include/ydb-cpp-sdk/library/operation_id/operation_id.h +++ b/include/ydb-cpp-sdk/library/operation_id/operation_id.h @@ -30,6 +30,7 @@ class TOperationId { RESTORE = 12, COMPACTION = 13, FULL_BACKUP = 14, + ANALYZE = 15, }; struct TData { diff --git a/src/api/grpc/ydb_topic_deferred_publish_v1.proto b/src/api/grpc/ydb_topic_deferred_publish_v1.proto new file mode 100644 index 00000000000..45bc9805985 --- /dev/null +++ b/src/api/grpc/ydb_topic_deferred_publish_v1.proto @@ -0,0 +1,25 @@ +syntax = "proto3"; +option cc_enable_arenas = true; + +package Ydb.Topic.DeferredPublish.V1; + +option java_package = "com.yandex.ydb.topic.deferredpublish.v1"; + +import "src/api/protos/ydb_topic_deferred_publish.proto"; + +service TopicDeferredPublishService { + // Starts a deferred topic publication. + rpc BeginPublication(DeferredPublish.BeginPublicationRequest) returns (DeferredPublish.BeginPublicationResponse); + + // Makes staged messages visible to readers. + rpc Publish(DeferredPublish.PublishRequest) returns (DeferredPublish.PublishResponse); + + // Discards staged messages without publishing. + rpc CancelPublication(DeferredPublish.CancelPublicationRequest) returns (DeferredPublish.CancelPublicationResponse); + + // Lists active publications. + rpc ListPublications(DeferredPublish.ListPublicationsRequest) returns (DeferredPublish.ListPublicationsResponse); + + // Returns metadata and destinations for an active publication. + rpc DescribePublication(DeferredPublish.DescribePublicationRequest) returns (DeferredPublish.DescribePublicationResponse); +} diff --git a/src/api/protos/draft/ydb_maintenance.proto b/src/api/protos/draft/ydb_maintenance.proto index 17f96b9847b..3b629b77052 100644 --- a/src/api/protos/draft/ydb_maintenance.proto +++ b/src/api/protos/draft/ydb_maintenance.proto @@ -98,6 +98,17 @@ message MaintenanceTaskOptions { bool dry_run = 4; // Priority of the task. Lower value indicates higher priority. int32 priority = 5 [(value) = "[-100; 100]"]; + // Maximum number of actions that may be in the PERFORMED state simultaneously + // for this task. CMS will not issue new locks beyond this limit; remaining + // actions stay in PENDING until clients call CompleteAction on already + // performed ones. 0 means unlimited (default behavior). + // Note: this option is honored only when actions can be granted + // independently — i.e. when the task contains more than one action_group + // or its single action_group contains exactly one action. For a task that + // consists of a single action_group with multiple actions (a "composite" + // group) actions within the group are granted atomically (all or nothing) + // and the cap is ignored — CMS will emit a warning in this case. + uint32 max_inflight_actions = 6; } // Used to describe the scope of a single action. diff --git a/src/api/protos/ydb_cms.proto b/src/api/protos/ydb_cms.proto index c72751609f5..fe90f5b1def 100644 --- a/src/api/protos/ydb_cms.proto +++ b/src/api/protos/ydb_cms.proto @@ -187,6 +187,7 @@ message GetDatabaseStatusResult { REMOVING = 3; PENDING_RESOURCES = 4; CONFIGURING = 5; + REMOVING_STORAGE_UNITS = 6; } // Full path to database's home dir. @@ -212,6 +213,9 @@ message GetDatabaseStatusResult { DatabaseQuotas database_quotas = 10; // Current scale recommender policies ScaleRecommenderPolicies scale_recommender_policies = 11; + // Outstanding problems related to database resources + // (e.g. storage pool allocation failures). + repeated Ydb.Issue.IssueMessage issues = 12; } // Change resources allocated for database. @@ -242,6 +246,8 @@ message AlterDatabaseRequest { map alter_attributes = 12; // Change scale recommender policies. ScaleRecommenderPolicies scale_recommender_policies = 13; + // Storage units to deallocate + repeated StorageUnits storage_units_to_remove = 14; } message AlterDatabaseResponse { diff --git a/src/api/protos/ydb_persqueue_v1.proto b/src/api/protos/ydb_persqueue_v1.proto index 6ba6fe7e42c..d95569c83e4 100644 --- a/src/api/protos/ydb_persqueue_v1.proto +++ b/src/api/protos/ydb_persqueue_v1.proto @@ -358,7 +358,7 @@ message StreamingReadClientMessage { int64 request_uncompressed_size = 1; } - // Signal for server that cient is ready to recive data for partition. + // Signal for server that cient is ready to recieve data for partition. message StartPartitionSessionResponse { // Partition session identifier of partition to start read. int64 partition_session_id = 1; diff --git a/src/api/protos/ydb_query.proto b/src/api/protos/ydb_query.proto index 4299f338fce..ecce71efb12 100644 --- a/src/api/protos/ydb_query.proto +++ b/src/api/protos/ydb_query.proto @@ -90,6 +90,9 @@ message SnapshotRWModeSettings { message ReadCommittedRWModeSettings { } +message StrictSerializableRWModeSettings { +} + message TransactionSettings { oneof tx_mode { SerializableModeSettings serializable_read_write = 1; @@ -98,6 +101,7 @@ message TransactionSettings { SnapshotModeSettings snapshot_read_only = 4; SnapshotRWModeSettings snapshot_read_write = 5; ReadCommittedRWModeSettings read_committed_read_write = 6; + StrictSerializableRWModeSettings strict_serializable_read_write = 7; } } diff --git a/src/api/protos/ydb_table.proto b/src/api/protos/ydb_table.proto index fac3b770311..b7ced0a0bd5 100644 --- a/src/api/protos/ydb_table.proto +++ b/src/api/protos/ydb_table.proto @@ -240,6 +240,7 @@ message FulltextIndexSettings { repeated ColumnAnalyzers columns = 2; } +// Without EnableCompactFulltextIndex feature flag (old version): // Simplest possible layout without support for relevance. // Uses a single flat inverted index table (indexImplTable). // Supports a single column only. @@ -263,6 +264,23 @@ message FulltextIndexSettings { // │ "quick" │ 1 │ // │ "quick" │ 2 │ // └──────────────┴────┘ +// +// With EnableCompactFulltextIndex feature flag (new version): +// Compact fulltext index format with segmented storage. +// Index table stores varint+delta-compressed binary segments (document ID lists) +// instead of individual document IDs. +// Example inverted index table (indexImplTable): +// ┌──────────────┬──────────────┬──────────────────┬─────────────┬───────────────┐ +// │ __ydb_token │ __ydb_max_id │ __ydb_generation | __ydb_added | __ydb_segment │ +// ├──────────────┼──────────────┼──────────────────┼─────────────┼───────────────┤ +// │ "The" │ 2^64-1 │ 2^64-1 │ true │ 1 2 │ +// │ "blue" │ 2^64-1 │ 2^64-1 │ true │ 2 │ +// │ "brown" │ 2^64-1 │ 2^64-1 │ true │ 1 │ +// │ "fox" │ 2^64-1 │ 2^64-1 │ true │ 1 │ +// │ "hare" │ 2^64-1 │ 2^64-1 │ true │ 2 │ +// │ "quick" │ 2^64-1 │ 2^64-1 │ true │ 1 2 │ +// └──────────────┴──────────────┴──────────────────┴─────────────┴───────────────┘ + message GlobalFulltextPlainIndex { GlobalIndexSettings settings = 1; FulltextIndexSettings fulltext_settings = 2; @@ -277,6 +295,7 @@ message GlobalFulltextPlainIndex { // │ 1 │ "A quick fox catches quick hare" │ // │ 2 │ "The quick blue hare" │ // └────┴───────────────────────────────────────┘ +// Without EnableCompactFulltextIndex feature flag (old version): // Example inverted index table with term frequencies (indexImplTable): // ┌──────────────┬────┬────────────┐ // │ __ydb_token │ id │ __ydb_freq │ @@ -291,6 +310,20 @@ message GlobalFulltextPlainIndex { // │ "quick" │ 1 │ 2 │ // │ "quick" │ 2 │ 1 │ // └──────────────┴────┴────────────┘ +// With EnableCompactFulltextIndex feature flag (new version): +// ┌──────────────┬──────────────┬──────────────────┬─────────────┬───────────────┐ +// │ __ydb_token │ __ydb_max_id │ __ydb_generation | __ydb_added | __ydb_segment │ +// ├──────────────┼──────────────┼──────────────────┼─────────────┼───────────────┤ +// │ "A" │ 2^64-1 │ 2^64-1 │ true │ 1:1 │ +// │ "The" │ 2^64-1 │ 2^64-1 │ true │ 2:1 │ +// │ "blue" │ 2^64-1 │ 2^64-1 │ true │ 2:1 │ +// │ "catches" │ 2^64-1 │ 2^64-1 │ true │ 1:1 │ +// │ "fox" │ 2^64-1 │ 2^64-1 │ true │ 1:1 │ +// │ "hare" │ 2^64-1 │ 2^64-1 │ true │ 1:1 2:1 │ +// │ "quick" │ 2^64-1 │ 2^64-1 │ true │ 1:2 2:1 │ +// └──────────────┴──────────────┴──────────────────┴─────────────┴───────────────┘ +// +// Other tables are the same for both compressed and uncompressed formats: // Term dictionary table (indexImplDictTable): // ┌──────────────┬────────────┐ // │ __ydb_token │ __ydb_freq │ @@ -416,6 +449,18 @@ message IndexBuildState { } } +// State of set column constraint operation +message SetColumnConstraintState { + enum State { + STATE_UNSPECIFIED = 0; + STATE_PREPARING = 1; + STATE_VALIDATING = 2; + STATE_APPLYING = 3; + STATE_DONE = 4; + STATE_CANCELLED = 5; + } +} + // Description of index building operation message IndexBuildDescription { string path = 1; @@ -1065,6 +1110,27 @@ message CompactMetadata { optional uint32 shards_done = 7; } +// State of an analyze operation +message AnalyzeState { + enum State { + STATE_UNSPECIFIED = 0; + STATE_ENQUEUED = 1; + STATE_IN_PROGRESS = 2; + STATE_DONE = 3; + STATE_CANCELLED = 4; + STATE_FAILED = 5; + } +} + +message AnalyzeMetadata { + + repeated string paths = 1; // All table paths covered by this analyze operation. + repeated string in_progress_paths = 2; // Tables currently being traversed + repeated string done_paths = 3; // Tables for which analysis is complete + optional AnalyzeState.State state = 4; + optional float progress = 5; +} + // Alter table with given path message AlterTableRequest { // Session identifier @@ -1284,6 +1350,9 @@ message SnapshotRWModeSettings { message ReadCommittedRWModeSettings { } +message StrictSerializableRWModeSettings { +} + message TransactionSettings { oneof tx_mode { SerializableModeSettings serializable_read_write = 1; @@ -1292,6 +1361,7 @@ message TransactionSettings { SnapshotModeSettings snapshot_read_only = 4; SnapshotRWModeSettings snapshot_read_write = 5; ReadCommittedRWModeSettings read_committed_read_write = 6; + StrictSerializableRWModeSettings strict_serializable_read_write = 7; } } diff --git a/src/api/protos/ydb_topic.proto b/src/api/protos/ydb_topic.proto index eba22480d7f..08aa58f643a 100644 --- a/src/api/protos/ydb_topic.proto +++ b/src/api/protos/ydb_topic.proto @@ -22,7 +22,8 @@ enum Codec { CODEC_GZIP = 2; CODEC_LZOP = 3; CODEC_ZSTD = 4; - reserved 5 to 9999; + CODEC_KAFKA_BATCH = 5; + reserved 6 to 9999; // User-defined codecs from 10000 to 19999 CODEC_CUSTOM = 10000; reserved 20000 to max; @@ -138,6 +139,9 @@ message StreamWriteMessage { // Client can only use compression codecs from this set to write messages to topic. // Otherwise session will be closed with BAD_REQUEST. SupportedCodecs supported_codecs = 4; + + // Indicates that server supports batching. + bool is_batching_supported = 5; } // Represents portion of client messages. @@ -148,10 +152,13 @@ message StreamWriteMessage { // See enum Codec above for values. int32 codec = 2; - optional TransactionIdentity tx = 3; + oneof write_context { + TransactionIdentity tx = 3; + DeferredPublishIdentity deferred_publish = 9; + } message MessageData { - // Message sequence number, provided by client for deduplication. + // Message sequence number, provided by client for deduplication. Max sequence number for batch messages. // Starts at 1 int64 seq_no = 1; // Creation timestamp. @@ -343,6 +350,9 @@ message StreamReadMessage { // Read only messages with 'written_at' value greater or equal than this timestamp. google.protobuf.Timestamp read_from = 4; } + + // Indicates that client supports batching. + bool is_batching_supported = 7; } // Handshake response. @@ -381,7 +391,7 @@ message StreamReadMessage { message MessageData { // Partition offset in partition that assigned for message. int64 offset = 1; //unique value for client side deduplication - Topic:Partition:Offset - // Sequence number that provided with message on write from client. + // Sequence number that provided with message on write from client. Max sequence number for batch messages. int64 seq_no = 2; // Timestamp of creation of message provided on write from client. google.protobuf.Timestamp created_at = 3; @@ -515,6 +525,10 @@ message StreamReadMessage { // All messages with offset less than commit_offset are processed by client. // Server will commit this position if this is not done yet. optional int64 commit_offset = 3; + + // If max_ofset is set, server will return data only up to max_offset (inclusive) + // If set and is less than read_offset, commit_offset or actual committed offset then server will send an error message (status != SUCCESS) and close stream. + optional int64 max_offset = 4; } // Command from server to stop and destroy concrete partition session. @@ -707,6 +721,15 @@ message TransactionIdentity { string session = 2; } +// Identity of a deferred topic publication for StreamWrite. +// int_publication_id is assigned by the server in BeginPublication. +// ext_publication_id is repeated from BeginPublication on each write batch. +message DeferredPublishIdentity { + uint64 int_publication_id = 1; + + string ext_publication_id = 2 [(Ydb.length).le = 2048]; +} + // Add offsets to transaction request sent from client to server. message UpdateOffsetsInTransactionRequest { Ydb.Operations.OperationParams operation_params = 1; diff --git a/src/api/protos/ydb_topic_deferred_publish.proto b/src/api/protos/ydb_topic_deferred_publish.proto new file mode 100644 index 00000000000..a7941fd16c9 --- /dev/null +++ b/src/api/protos/ydb_topic_deferred_publish.proto @@ -0,0 +1,106 @@ +syntax = "proto3"; +option cc_enable_arenas = true; + +import "src/api/protos/annotations/validation.proto"; +import "src/api/protos/ydb_operation.proto"; + +import "google/protobuf/timestamp.proto"; + +package Ydb.Topic.DeferredPublish; + +option java_package = "com.yandex.ydb.topic.deferredpublish"; + +message BeginPublicationRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Client-defined publication name. Unique among active publications in the database. + string ext_publication_id = 2 [(Ydb.length).le = 2048]; + + // Optional writer identity stored in the registry and used as a filter in ListPublications. + optional string writer_identity = 3 [(Ydb.length).le = 2048]; +} + +message BeginPublicationResponse { + Ydb.Operations.Operation operation = 1; +} + +message BeginPublicationResult { + uint64 int_publication_id = 1; +} + +message PublishRequest { + Ydb.Operations.OperationParams operation_params = 1; + + uint64 int_publication_id = 2; +} + +message PublishResponse { + Ydb.Operations.Operation operation = 1; +} + +message PublishResult { +} + +message CancelPublicationRequest { + Ydb.Operations.OperationParams operation_params = 1; + + uint64 int_publication_id = 2; +} + +message CancelPublicationResponse { + Ydb.Operations.Operation operation = 1; +} + +message CancelPublicationResult { +} + +message ListPublicationsRequest { + Ydb.Operations.OperationParams operation_params = 1; + + // Optional filter by writer identity from BeginPublication. + optional string writer_identity = 2 [(Ydb.length).le = 2048]; +} + +message ListPublicationsResponse { + Ydb.Operations.Operation operation = 1; +} + +message PublicationSummary { + uint64 int_publication_id = 1; + + string ext_publication_id = 2 [(Ydb.length).le = 2048]; + + optional string writer_identity = 3 [(Ydb.length).le = 2048]; +} + +message ListPublicationsResult { + repeated PublicationSummary publications = 1; +} + +message DescribePublicationRequest { + Ydb.Operations.OperationParams operation_params = 1; + + uint64 int_publication_id = 2; +} + +message DescribePublicationResponse { + Ydb.Operations.Operation operation = 1; +} + +message DescribePublicationResult { + string ext_publication_id = 1; + + optional string writer_identity = 2; + + google.protobuf.Timestamp created_at = 3; + + optional string created_by = 4; + + repeated Destination destinations = 5; + + message Destination { + string topic_path = 1; + + repeated int64 partition_ids = 2; + } +} diff --git a/src/client/CMakeLists.txt b/src/client/CMakeLists.txt index cc888987a05..e9b85f044ea 100644 --- a/src/client/CMakeLists.txt +++ b/src/client/CMakeLists.txt @@ -22,6 +22,7 @@ add_subdirectory(query) add_subdirectory(rate_limiter) add_subdirectory(resources) add_subdirectory(result) +add_subdirectory(row_ranges) add_subdirectory(scheme) add_subdirectory(secret) add_subdirectory(ss_tasks) diff --git a/src/client/federated_topic/impl/federated_deferred_commit.cpp b/src/client/federated_topic/impl/federated_deferred_commit.cpp index dc402c05f55..15e70ebf3c2 100644 --- a/src/client/federated_topic/impl/federated_deferred_commit.cpp +++ b/src/client/federated_topic/impl/federated_deferred_commit.cpp @@ -8,10 +8,10 @@ namespace NYdb::inline V3::NFederatedTopic { std::pair GetMessageOffsetRange(const TReadSessionEvent::TDataReceivedEvent& dataReceivedEvent, ui64 index) { if (dataReceivedEvent.HasCompressedMessages()) { const auto& msg = dataReceivedEvent.GetCompressedMessages()[index]; - return {msg.GetOffset(), msg.GetOffset() + 1}; + return {msg.GetOffset(), msg.GetOffset() + msg.GetLogicalMessageCount()}; } const auto& msg = dataReceivedEvent.GetMessages()[index]; - return {msg.GetOffset(), msg.GetOffset() + 1}; + return {msg.GetOffset(), msg.GetOffset() + msg.GetLogicalMessageCount()}; } diff --git a/src/client/federated_topic/impl/federated_read_session_event.cpp b/src/client/federated_topic/impl/federated_read_session_event.cpp index 658ae75d859..b812697b407 100644 --- a/src/client/federated_topic/impl/federated_read_session_event.cpp +++ b/src/client/federated_topic/impl/federated_read_session_event.cpp @@ -137,6 +137,20 @@ void TPrintable::DebugString(TStringBuilder& ret, bool print namespace NYdb::inline V3::NFederatedTopic { +namespace { + +void AddOffsetRange(std::vector>& offsetRanges, const NTopic::TReadSessionEvent::TDataReceivedEvent::TMessageBase& msg) { + const auto from = msg.GetOffset(); + const auto to = from + msg.GetLogicalMessageCount(); + if (offsetRanges.empty() || offsetRanges.back().second != from) { + offsetRanges.emplace_back(from, to); + } else { + offsetRanges.back().second = to; + } +} + +} // namespace + //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// // NFederatedTopic::TReadSessionEvent::TDataReceivedEvent @@ -146,10 +160,12 @@ TReadSessionEvent::TDataReceivedEvent::TDataReceivedEvent(NTopic::TReadSessionEv { if (event.HasCompressedMessages()) { for (auto& msg : event.GetCompressedMessages()) { + AddOffsetRange(OffsetRanges, msg); CompressedMessages.emplace_back(std::move(msg), federatedPartitionSession); } } else { for (auto& msg : event.GetMessages()) { + AddOffsetRange(OffsetRanges, msg); Messages.emplace_back(std::move(msg), federatedPartitionSession); } } diff --git a/src/client/federated_topic/impl/federation_observer.cpp b/src/client/federated_topic/impl/federation_observer.cpp index 437e88c4ee7..14f2b58e22a 100644 --- a/src/client/federated_topic/impl/federation_observer.cpp +++ b/src/client/federated_topic/impl/federation_observer.cpp @@ -62,14 +62,30 @@ Ydb::FederationDiscovery::ListFederationDatabasesRequest TFederatedDbObserverImp } void TFederatedDbObserverImpl::RunFederationDiscoveryImpl() { - Y_ABORT_UNLESS(Lock.IsLocked()); + // NOTE: must be called WITHOUT Lock held. + // RunDeferred may invoke the response callback synchronously (e.g. when the + // IAM token is not ready yet), which in turn calls OnFederationDiscovery that + // tries to acquire Lock. Holding Lock here would cause a spin-deadlock on + // the single gRPC event-loop thread. + Y_ABORT_UNLESS(!Lock.IsLocked()); - FederationDiscoveryDelayContext = Connections_->CreateContext(); - if (!FederationDiscoveryDelayContext) { - Stopping = true; - // TODO log DRIVER_IS_STOPPING_DESCRIPTION - return; + NYdbGrpc::IQueueClientContextPtr ctx; + { + std::lock_guard guard(Lock); + if (Stopping) { + return; + } + ctx = Connections_->CreateContext(); + if (!ctx) { + Stopping = true; + FederationDiscoveryDelayContext = nullptr; // release any previously-held context + // TODO log DRIVER_IS_STOPPING_DESCRIPTION + return; + } + FederationDiscoveryDelayContext = ctx; } + // Lock is released here — before RunDeferred — to prevent the deadlock + // described above. auto extractor = [selfCtx = SelfContext] (google::protobuf::Any* any, TPlainStatus status) mutable { @@ -94,7 +110,7 @@ void TFederatedDbObserverImpl::RunFederationDiscoveryImpl() { DbDriverState_, {}, // no polling unready operations, so no need in delay parameter settings, - FederationDiscoveryDelayContext); + std::move(ctx)); } void TFederatedDbObserverImpl::ScheduleFederationDiscoveryImpl(TDuration delay) { @@ -102,10 +118,17 @@ void TFederatedDbObserverImpl::ScheduleFederationDiscoveryImpl(TDuration delay) auto cb = [selfCtx = SelfContext](bool ok) { if (ok) { if (auto self = selfCtx->LockShared()) { - std::lock_guard guard(self->Lock); - if (self->Stopping) { - return; + { + std::lock_guard guard(self->Lock); + if (self->Stopping) { + return; + } } + // Lock is intentionally released before RunFederationDiscoveryImpl: + // that function must not be called with Lock held because RunDeferred + // can invoke its callback synchronously (e.g. on auth errors), which + // would call OnFederationDiscovery → try to re-acquire the same + // non-recursive TSpinLock → spin-deadlock. self->RunFederationDiscoveryImpl(); } } diff --git a/src/client/federated_topic/ut/federation_observer_deadlock_ut.cpp b/src/client/federated_topic/ut/federation_observer_deadlock_ut.cpp new file mode 100644 index 00000000000..28d4f0f5031 --- /dev/null +++ b/src/client/federated_topic/ut/federation_observer_deadlock_ut.cpp @@ -0,0 +1,88 @@ +#include +#include +#include +#include + +#include + +#include +#include + +#include + +#include + +namespace NYdb::NFederatedTopic::NTests { + +namespace { + +class TNotReadyTokenCredentialsProvider final : public ICredentialsProvider { +public: + std::string GetAuthInfo() const override { + throw TAuthenticationError("IAM-token not ready yet"); + } + + bool IsValid() const override { + return true; + } +}; + +class TNotReadyTokenCredentialsProviderFactory final : public ICredentialsProviderFactory { +public: + TCredentialsProviderPtr CreateProvider() const override { + return std::make_shared(); + } + + TCredentialsProviderPtr CreateProvider(std::weak_ptr facility) const override { + Y_UNUSED(facility); + return CreateProvider(); + } +}; + +std::unique_ptr StartFdsGrpcServer(TFederationDiscoveryServiceMock& service, int* port) { + grpc::ServerBuilder builder; + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), port); + builder.RegisterService(&service); + return builder.BuildAndStart(); +} + +} // namespace + +Y_UNIT_TEST_SUITE(FederationObserverDeadlock) { + +// Regression for YQ-5434: RunFederationDiscoveryImpl must not hold Lock across RunDeferred. +// When GetAuthInfo() fails synchronously, RunDeferred invokes the response callback inline, +// which calls OnFederationDiscovery and tries to re-acquire the same non-recursive TSpinLock. +Y_UNIT_TEST(SyncAuthErrorDuringDiscoveryDoesNotDeadlock) { + TFederationDiscoveryServiceMock fdsMock; + + int port = 0; + auto grpcServer = StartFdsGrpcServer(fdsMock, &port); + UNIT_ASSERT(grpcServer); + UNIT_ASSERT(port > 0); + + auto credentialsFactory = std::make_shared(); + + TDriverConfig driverConfig; + driverConfig + .SetEndpoint(TStringBuilder() << "127.0.0.1:" << port) + .SetDatabase("/Root") + .SetCredentialsProviderFactory(credentialsFactory); + + TDriver driver(driverConfig); + + TFederatedTopicClientSettings clientSettings; + clientSettings + .CredentialsProviderFactory(credentialsFactory) + .RetryPolicy(NTopic::IRetryPolicy::GetNoRetryPolicy()); + + TFederatedTopicClient client(driver, clientSettings); + + auto clusterInfoFuture = client.GetAllClusterInfo(); + UNIT_ASSERT(clusterInfoFuture.Wait(TDuration::Seconds(5))); + UNIT_ASSERT_VALUES_EQUAL(clusterInfoFuture.GetValue().size(), 0u); +} + +} // Y_UNIT_TEST_SUITE(FederationObserverDeadlock) + +} // namespace NYdb::NFederatedTopic::NTests diff --git a/src/client/federated_topic/ut/simple_blocking_write_session_ut.cpp b/src/client/federated_topic/ut/simple_blocking_write_session_ut.cpp index 51260f74b12..baec2a04dc2 100644 --- a/src/client/federated_topic/ut/simple_blocking_write_session_ut.cpp +++ b/src/client/federated_topic/ut/simple_blocking_write_session_ut.cpp @@ -1,6 +1,8 @@ #include #include +#include + #include #include @@ -19,6 +21,7 @@ class TSimpleBlockingWriteSessionTestFixture { testName, false, ::NPersQueue::TTestServer::LOGGED_SERVICES, NActors::NLog::PRI_DEBUG, 1)) , ThreadPool(CreateThreadPool(2)) { + NKikimr::NPQ::NTest::EnableTopicBatching(Setup->GetServer().ServerSettings); Setup->Start(true, true); FdsMock.Port = Setup->GetGrpcPort(); @@ -78,11 +81,19 @@ class TSimpleBlockingWriteSessionTestFixture { } std::shared_ptr CreateReadSession() { + return CreateReadSession(DefaultReadSettings()); + } + + TFederatedReadSessionSettings DefaultReadSettings() { TFederatedReadSessionSettings settings; settings .ConsumerName("shared/user") .MaxMemoryUsageBytes(1_MB) .AppendTopics(std::string(Setup->GetTestTopicPath())); + return settings; + } + + std::shared_ptr CreateReadSession(const TFederatedReadSessionSettings& settings) { return TopicClient->CreateReadSession(settings); } @@ -116,7 +127,13 @@ class TSimpleBlockingWriteSessionTestFixture { // Creates a read session with single-database FDS response (avoids partition competition) std::shared_ptr CreateReadSessionWithSingleDatabaseFdsResponse() { - auto session = CreateReadSession(); + return CreateReadSessionWithSingleDatabaseFdsResponse(DefaultReadSettings()); + } + + std::shared_ptr CreateReadSessionWithSingleDatabaseFdsResponse( + const TFederatedReadSessionSettings& settings + ) { + auto session = CreateReadSession(settings); auto fdsRequest = FdsMock.WaitNextPendingRequest(); fdsRequest.Result.SetValue(FdsMock.ComposeOkResultSingleDatabase()); @@ -175,6 +192,22 @@ class TSimpleBlockingWriteSessionTestFixture { TString GetTestTopic() const { return Setup->GetTestTopic(); } TFederationDiscoveryServiceMock& GetFdsMock() { return FdsMock; } + void WriteKafkaBatchMessages( + const TString& producerId, + size_t dataSize, + ui64 maxBatchMessageCount, + const TVector>& writes + ) { + NTopic::TTopicClient client(Setup->GetDriver()); + NKikimr::NPQ::NTest::WriteKafkaBatchMessages( + client, + Setup->GetTestTopicPath(), + producerId, + dataSize, + maxBatchMessageCount, + writes); + } + // Helper to verify that messages were written correctly by reading them back void VerifyMessagesWritten(const TVector& expectedMessages) { // Create read session with single database FDS response @@ -397,6 +430,101 @@ Y_UNIT_TEST_SUITE(SimpleBlockingFederatedWriteSession) { "Write should fail after session is closed"); } + Y_UNIT_TEST(DeferredCommitKafkaBatches) { + TSimpleBlockingWriteSessionTestFixture fixture(TEST_CASE_NAME); + + const TString producerId = "federated-kafka-batch-producer"; + constexpr size_t dataSize = 16; + constexpr ui64 maxBatchMessageCount = 5; + constexpr ui64 totalMessages = 8; + + fixture.WriteKafkaBatchMessages( + producerId, + dataSize, + maxBatchMessageCount, + { + {1, 5, 'a'}, + {6, 3, 'b'}, + }); + + auto readSettings = fixture.DefaultReadSettings(); + readSettings.Decompress(false); + + { + auto readSession = fixture.CreateReadSessionWithSingleDatabaseFdsResponse(readSettings); + fixture.GetFdsMock().SetAutoRespondSingleDatabase(true); + + TDeferredCommit deferredCommit; + TFederatedPartitionSession::TPtr partitionSession; + ui64 logicalMessagesRead = 0; + ui64 compressedBatchesRead = 0; + bool sawMultiMessageBatch = false; + const TInstant deadline = TInstant::Now() + TDuration::Seconds(30); + while (logicalMessagesRead < totalMessages && TInstant::Now() < deadline) { + UNIT_ASSERT_C(readSession->WaitEvent().Wait(TDuration::Seconds(5)), + "Read session event timeout"); + auto event = readSession->GetEvent(false); + UNIT_ASSERT(event.has_value()); + + if (auto* start = std::get_if(&*event)) { + partitionSession = start->GetFederatedPartitionSession(); + start->Confirm(); + continue; + } + if (auto* stop = std::get_if(&*event)) { + stop->Confirm(); + continue; + } + if (auto* data = std::get_if(&*event)) { + UNIT_ASSERT(data->HasCompressedMessages()); + for (const auto& message : data->GetCompressedMessages()) { + UNIT_ASSERT_VALUES_EQUAL(message.GetOffset(), logicalMessagesRead); + logicalMessagesRead += message.GetLogicalMessageCount(); + sawMultiMessageBatch |= message.GetLogicalMessageCount() > 1; + ++compressedBatchesRead; + } + deferredCommit.Add(*data); + } + } + + UNIT_ASSERT_VALUES_EQUAL(logicalMessagesRead, totalMessages); + UNIT_ASSERT_VALUES_EQUAL(compressedBatchesRead, 2u); + UNIT_ASSERT(sawMultiMessageBatch); + + deferredCommit.Commit(); + + UNIT_ASSERT(partitionSession); + partitionSession->RequestStatus(); + + bool gotStatus = false; + const TInstant statusDeadline = TInstant::Now() + TDuration::Seconds(30); + while (!gotStatus && TInstant::Now() < statusDeadline) { + if (!readSession->WaitEvent().Wait(TDuration::Seconds(1))) { + partitionSession->RequestStatus(); + continue; + } + auto event = readSession->GetEvent(false); + UNIT_ASSERT(event.has_value()); + + if (auto* status = std::get_if(&*event)) { + if (status->GetCommittedOffset() == totalMessages) { + gotStatus = true; + } else { + partitionSession->RequestStatus(); + } + continue; + } + if (auto* stop = std::get_if(&*event)) { + stop->Confirm(); + } + } + + UNIT_ASSERT_C(gotStatus, "Did not observe committed offset " << totalMessages); + + UNIT_ASSERT(readSession->Close(TDuration::Seconds(5))); + } + } + } } // namespace NYdb::NFederatedTopic::NTests diff --git a/src/client/impl/endpoints/endpoints.cpp b/src/client/impl/endpoints/endpoints.cpp index 670874a45e3..d6d12c3c558 100644 --- a/src/client/impl/endpoints/endpoints.cpp +++ b/src/client/impl/endpoints/endpoints.cpp @@ -2,7 +2,7 @@ #include #include - +#include #include #include @@ -52,6 +52,8 @@ class TEndpointElectorSafe::TObjRegistry : public IObjRegistryHandle { //////////////////////////////////////////////////////////////////////////////// +constexpr std::int32_t PessimizedPriority = std::numeric_limits::max(); + // Returns index of last resord with same priority or -1 in case of empty input static std::int32_t GetBestK(const std::vector& records) { if (records.empty()) { @@ -157,27 +159,68 @@ TEndpointRecord TEndpointElectorSafe::GetEndpoint(const TEndpointKey& preferredE return Records_[idx]; } } +bool TEndpointElectorSafe::PessimizeEndpointUnlocked(const std::string& endpoint) { + const auto recordIt = std::ranges::find_if(Records_, [&](const TEndpointRecord& r) { + return r.Endpoint == endpoint && r.Priority != PessimizedPriority; + }); + if (recordIt == Records_.end()) { + return false; + } + + const auto nodeId = recordIt->NodeId; + recordIt->Priority = PessimizedPriority; + + if (auto it = KnownEndpoints_.find(endpoint); it != KnownEndpoints_.end()) { + it->second.Priority = PessimizedPriority; + } + + if (auto it = KnownEndpointsByNodeId_.find(nodeId); it != KnownEndpointsByNodeId_.end() + && it->second.Record.Endpoint == endpoint) + { + it->second.Record.Priority = PessimizedPriority; + } + + const int pessimizationRatio = PessimizationRatio_.load(); + const auto newRatio = (pessimizationRatio * Records_.size() + 100) / Records_.size(); + PessimizationRatio_.store(newRatio); + PessimizationRatioGauge_.SetValue(newRatio); + EndpointActiveGauge_.Dec(); + + return true; +} -// TODO: Suboptimal, but should not be used often void TEndpointElectorSafe::PessimizeEndpoint(const std::string& endpoint) { std::unique_lock guard(Mutex_); - for (auto& r : Records_) { - if (r.Endpoint == endpoint && r.Priority != std::numeric_limits::max()) { - int pessimizationRatio = PessimizationRatio_.load(); - auto newRatio = (pessimizationRatio * Records_.size() + 100) / Records_.size(); - PessimizationRatio_.store(newRatio); - PessimizationRatioGauge_.SetValue(newRatio); - EndpointActiveGauge_.Dec(); - r.Priority = std::numeric_limits::max(); - - auto it = KnownEndpoints_.find(endpoint); - if (it != KnownEndpoints_.end()) { - it->second.Priority = std::numeric_limits::max(); - } - } + if (PessimizeEndpointUnlocked(endpoint)) { + Sort(Records_.begin(), Records_.end()); + BestK_ = GetBestK(Records_); + } +} + +void TEndpointElectorSafe::PessimizeNode(const std::uint64_t nodeId) { + if (nodeId == 0) { + return; + } + + std::unique_lock guard(Mutex_); + + auto endpoints = Records_ + | std::views::filter([nodeId](const TEndpointRecord& r) { + return r.NodeId == nodeId && r.Priority != PessimizedPriority; + }) + | std::views::transform([](const TEndpointRecord& r) { + return r.Endpoint; + }); + + bool changed = false; + for (const auto& endpoint : endpoints) { + changed |= PessimizeEndpointUnlocked(endpoint); + } + + if (changed) { + Sort(Records_.begin(), Records_.end()); + BestK_ = GetBestK(Records_); } - Sort(Records_.begin(), Records_.end()); - BestK_ = GetBestK(Records_); } // % of endpoints which was pessimized diff --git a/src/client/impl/endpoints/endpoints.h b/src/client/impl/endpoints/endpoints.h index 0e6d288a3fb..cb7508fa8ac 100644 --- a/src/client/impl/endpoints/endpoints.h +++ b/src/client/impl/endpoints/endpoints.h @@ -100,6 +100,9 @@ class TEndpointElectorSafe { // Move endpoint to the end void PessimizeEndpoint(const std::string& endpoint); + // Move all endpoints of a certain node to the end + void PessimizeNode(const std::uint64_t nodeId); + // Returns % of pessimized endpoints int GetPessimizationRatio() const; @@ -113,6 +116,9 @@ class TEndpointElectorSafe { class TObjRegistry; private: + // Move endpoint to the end + bool PessimizeEndpointUnlocked(const std::string& endpoint); + using TTaggedObjRegistry = std::unordered_map>; struct TKnownEndpoint { diff --git a/src/client/impl/internal/db_driver_state/endpoint_pool.cpp b/src/client/impl/internal/db_driver_state/endpoint_pool.cpp index f122a453e87..0c8b756eaf6 100644 --- a/src/client/impl/internal/db_driver_state/endpoint_pool.cpp +++ b/src/client/impl/internal/db_driver_state/endpoint_pool.cpp @@ -140,6 +140,10 @@ void TEndpointPool::BanEndpoint(const std::string& endpoint) { Elector_.PessimizeEndpoint(endpoint); } +void TEndpointPool::BanNodeId(std::uint64_t nodeId) { + Elector_.PessimizeNode(nodeId); +} + int TEndpointPool::GetPessimizationRatio() { return Elector_.GetPessimizationRatio(); } diff --git a/src/client/impl/internal/db_driver_state/endpoint_pool.h b/src/client/impl/internal/db_driver_state/endpoint_pool.h index b5345933372..2b11e98aa33 100644 --- a/src/client/impl/internal/db_driver_state/endpoint_pool.h +++ b/src/client/impl/internal/db_driver_state/endpoint_pool.h @@ -35,6 +35,7 @@ class TEndpointPool { TEndpointRecord GetEndpoint(const TEndpointKey& preferredEndpoint, bool onlyPreferred = false) const; TDuration TimeSinceLastUpdate() const; void BanEndpoint(const std::string& endpoint); + void BanNodeId(std::uint64_t nodeId); int GetPessimizationRatio(); bool LinkObjToEndpoint(const TEndpointKey& endpoint, TEndpointObj* obj, const void* tag); void ForEachEndpoint(const TEndpointElectorSafe::THandleCb& cb, const void* tag) const; diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.cpp b/src/client/impl/internal/grpc_connections/grpc_connections.cpp index 6d6b76a2132..f739bf20aa6 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.cpp +++ b/src/client/impl/internal/grpc_connections/grpc_connections.cpp @@ -26,7 +26,9 @@ std::string GetAuthInfo(TDbDriverStatePtr p) { } return token; } catch (const TAuthenticationError& e) { - throw e; + throw; + } catch (const TYdbException& e) { + throw; } catch (const std::exception& e) { throw TAuthenticationError(TStringBuilder() << "Can't get Authentication info from CredentialsProvider. " << e.what()); } diff --git a/src/client/impl/internal/grpc_connections/grpc_connections.h b/src/client/impl/internal/grpc_connections/grpc_connections.h index 0fa77f5ed69..329cec719e1 100644 --- a/src/client/impl/internal/grpc_connections/grpc_connections.h +++ b/src/client/impl/internal/grpc_connections/grpc_connections.h @@ -260,10 +260,10 @@ class TGRpcConnectionsImpl try { meta = MakeCallMeta(requestSettings, dbState); - } catch (const TAuthenticationError& e) { + } catch (const TYdbException& e) { userResponseCb( nullptr, - TPlainStatus(EStatus::CLIENT_UNAUTHENTICATED, e.what()) + TPlainStatus(dynamic_cast(&e) ? EStatus::CLIENT_UNAUTHENTICATED : EStatus::UNAVAILABLE, e.what()) ); return; } @@ -475,9 +475,9 @@ class TGRpcConnectionsImpl TCallMeta meta; try { meta = MakeCallMeta(requestSettings, dbState); - } catch (const TAuthenticationError& e) { + } catch (const TYdbException& e) { responseCb( - TPlainStatus(EStatus::CLIENT_UNAUTHENTICATED, e.what()), + TPlainStatus(dynamic_cast(&e) ? EStatus::CLIENT_UNAUTHENTICATED : EStatus::UNAVAILABLE, e.what()), nullptr ); return; @@ -556,9 +556,9 @@ class TGRpcConnectionsImpl TCallMeta meta; try { meta = MakeCallMeta(requestSettings, dbState); - } catch (const TAuthenticationError& e) { + } catch (const TYdbException& e) { connectedCallback( - TPlainStatus(EStatus::CLIENT_UNAUTHENTICATED, e.what()), + TPlainStatus(dynamic_cast(&e) ? EStatus::CLIENT_UNAUTHENTICATED : EStatus::UNAVAILABLE, e.what()), nullptr ); return; diff --git a/src/client/impl/internal/retry/CMakeLists.txt b/src/client/impl/internal/retry/CMakeLists.txt index d8d58530000..3e3bfdf0102 100644 --- a/src/client/impl/internal/retry/CMakeLists.txt +++ b/src/client/impl/internal/retry/CMakeLists.txt @@ -8,6 +8,7 @@ target_link_libraries(impl-internal-retry PUBLIC target_sources(impl-internal-retry PRIVATE retry.cpp + bulk_upsert_retry_state.cpp ) _ydb_sdk_install_targets(TARGETS impl-internal-retry) diff --git a/src/client/impl/internal/retry/bulk_upsert_retry_state.cpp b/src/client/impl/internal/retry/bulk_upsert_retry_state.cpp new file mode 100644 index 00000000000..5af875ad20d --- /dev/null +++ b/src/client/impl/internal/retry/bulk_upsert_retry_state.cpp @@ -0,0 +1,29 @@ +#include "bulk_upsert_retry_state.h" + +#include + +namespace NYdb::inline V3::NRetry { + +TBulkUpsertRetryState::TBulkUpsertRetryState(TRetryOperationSettings settings) + : Settings_(std::move(settings)) +{} + +const TRetryOperationSettings& TBulkUpsertRetryState::Settings() const { + return Settings_; +} + +bool TBulkUpsertRetryState::HasBackup() const { + return Backup_.has_value(); +} + +void TBulkUpsertRetryState::CreateBackup(const Ydb::Table::BulkUpsertRequest& request) { + const auto& rows = request.rows(); + Backup_.emplace(NYdb::TType(rows.type()), rows.value()); +} + +NYdb::TValue TBulkUpsertRetryState::GetBackupCopy() const { + Y_ABORT_UNLESS(Backup_.has_value()); + return NYdb::TValue(*Backup_); +} + +} // namespace NYdb::NRetry diff --git a/src/client/impl/internal/retry/bulk_upsert_retry_state.h b/src/client/impl/internal/retry/bulk_upsert_retry_state.h new file mode 100644 index 00000000000..5b14bf65190 --- /dev/null +++ b/src/client/impl/internal/retry/bulk_upsert_retry_state.h @@ -0,0 +1,30 @@ +#pragma once + +#include +#include +#include + +#include + +namespace Ydb::Table { +class BulkUpsertRequest; +} + +namespace NYdb::inline V3::NRetry { + +class TBulkUpsertRetryState { +public: + explicit TBulkUpsertRetryState(TRetryOperationSettings settings); + + const TRetryOperationSettings& Settings() const; + + bool HasBackup() const; + void CreateBackup(const Ydb::Table::BulkUpsertRequest& request); + NYdb::TValue GetBackupCopy() const; + +private: + TRetryOperationSettings Settings_; + std::optional Backup_; +}; + +} // namespace NYdb::NRetry diff --git a/src/client/impl/internal/retry/retry.h b/src/client/impl/internal/retry/retry.h index 2a25f0bfad7..9a6106dd24a 100644 --- a/src/client/impl/internal/retry/retry.h +++ b/src/client/impl/internal/retry/retry.h @@ -14,6 +14,7 @@ #include #include +#include namespace NYdb::inline V3 { class IClientImplCommon; @@ -32,6 +33,25 @@ enum class NextStep { Finish, }; +inline bool ShouldRetryStatus(EStatus status, const TRetryOperationSettings& settings) { + switch (status) { + case EStatus::ABORTED: + case EStatus::OVERLOADED: + case EStatus::CLIENT_RESOURCE_EXHAUSTED: + case EStatus::UNAVAILABLE: + case EStatus::BAD_SESSION: + case EStatus::SESSION_BUSY: + return true; + case EStatus::NOT_FOUND: + return settings.RetryNotFound_; + case EStatus::UNDETERMINED: + case EStatus::TRANSPORT_UNAVAILABLE: + return settings.Idempotent_; + default: + return settings.RetryUndefined_; + } +} + class TRetryContextBase : TNonCopyable { protected: TRetryOperationSettings Settings_; @@ -116,6 +136,20 @@ class TRetryContextBase : TNonCopyable { } }; +template +TStatusType MakeRetryResultFromStatus(TStatus&& status) { + return TStatusType(TStatus(std::move(status))); +} + +template +TStatusType InvokeWithRangeErrorCatch(F&& f) { + try { + return f(); + } catch (const NStatusHelpers::TYdbRangeErrorException& e) { + return MakeRetryResultFromStatus(TStatus(e.GetStatus())); + } +} + template class TRetryDeadlineHelper { public: @@ -124,4 +158,37 @@ class TRetryDeadlineHelper { } }; +template +class TInRetryOperationContextClientGuard { +public: + explicit TInRetryOperationContextClientGuard(TClient& client) + : Client_(client) + , Previous_(client.GetInRetryOperationContext()) + { + Client_.SetInRetryOperationContext(true); + } + + ~TInRetryOperationContextClientGuard() { + Client_.SetInRetryOperationContext(Previous_); + } + +private: + TClient& Client_; + bool Previous_; +}; + +template +const TStatus& GetRetryStatus(const TStatusType& status) { + if constexpr (std::is_base_of_v>) { + return status; + } else { + return status.Status(); + } +} + +template +EStatus GetRetryStatusCode(const TStatusType& status) { + return GetRetryStatus(status).GetStatus(); +} + } // namespace NYdb::NRetry diff --git a/src/client/impl/internal/retry/retry_async.h b/src/client/impl/internal/retry/retry_async.h index acb1987b26b..25ca95d759f 100644 --- a/src/client/impl/internal/retry/retry_async.h +++ b/src/client/impl/internal/retry/retry_async.h @@ -40,7 +40,7 @@ class TRetryContext : public TThrRefBase, public TRetryContextBase { auto value = f.GetValue(); if (self->ParentSpan_) { self->ParentSpan_->SetRetryCount(self->RetryNumber_); - self->ParentSpan_->End(value.GetStatus()); + self->ParentSpan_->End(GetRetryStatusCode(value)); } return value; } catch (...) { @@ -97,12 +97,13 @@ class TRetryContext : public TThrRefBase, public TRetryContextBase { } static void HandleStatusAsync(TPtr self, const TStatusType& status) { - self->EndAttemptSpan(status.GetStatus()); - auto nextStep = self->GetNextStep(status); + const TStatus& retryStatus = GetRetryStatus(status); + self->EndAttemptSpan(retryStatus.GetStatus()); + auto nextStep = self->GetNextStep(retryStatus); if (nextStep != NextStep::Finish) { self->RetryNumber_++; - self->Client_.Impl_->CollectRetryStatAsync(status.GetStatus()); - self->LogRetry(status); + self->Client_.Impl_->CollectRetryStatAsync(retryStatus.GetStatus()); + self->LogRetry(retryStatus); } switch (nextStep) { case NextStep::RetryImmediately: @@ -118,16 +119,24 @@ class TRetryContext : public TThrRefBase, public TRetryContextBase { } static void DoRunOperation(TPtr self) { - self->RunOperation().Subscribe( - [self](const TAsyncStatusType& result) { - [[maybe_unused]] auto attemptScope = self->ActivateAttemptSpan(); - try { - HandleStatusAsync(self, result.GetValue()); - } catch (...) { - HandleExceptionAsync(self, std::current_exception()); + try { + self->RunOperation().Subscribe( + [self](const TAsyncStatusType& result) { + [[maybe_unused]] auto attemptScope = self->ActivateAttemptSpan(); + try { + HandleStatusAsync(self, result.GetValue()); + } catch (const NStatusHelpers::TYdbRangeErrorException& e) { + HandleStatusAsync(self, MakeRetryResultFromStatus(TStatus(e.GetStatus()))); + } catch (...) { + HandleExceptionAsync(self, std::current_exception()); + } } - } - ); + ); + } catch (const NStatusHelpers::TYdbRangeErrorException& e) { + HandleStatusAsync(self, MakeRetryResultFromStatus(TStatus(e.GetStatus()))); + } catch (...) { + HandleExceptionAsync(self, std::current_exception()); + } } protected: @@ -175,6 +184,7 @@ class TRetryWithoutSession : public TRetryContext { protected: TAsyncStatusType RunOperation() override { + TInRetryOperationContextClientGuard guard(this->Client_); if constexpr (TFunctionArgs::Length == 1) { return Operation_(this->Client_); } else { @@ -217,7 +227,8 @@ class TRetryWithSession : public TRetryContext, publi try { auto& result = resultFuture.GetValue(); if (!result.IsSuccess()) { - return TRetryContextAsync::HandleStatusAsync(self, TStatusType(TStatus(result))); + return TRetryContextAsync::HandleStatusAsync( + self, MakeRetryResultFromStatus(TStatus(result))); } self->Session_ = result.GetSession(); @@ -239,6 +250,7 @@ class TRetryWithSession : public TRetryContext, publi } TAsyncStatusType RunOperation() override { + TInRetryOperationContextClientGuard guard(this->Client_); if constexpr (TFunctionArgs::Length == 1) { return Operation_(this->Session_.value()); } else { diff --git a/src/client/impl/internal/retry/retry_settings.h b/src/client/impl/internal/retry/retry_settings.h new file mode 100644 index 00000000000..ebe20aceb4f --- /dev/null +++ b/src/client/impl/internal/retry/retry_settings.h @@ -0,0 +1,78 @@ +#pragma once + +#include +#include + +#include +#include + +namespace NYdb::inline V3::NRetry { + +enum class ERetryIdempotentDefault { + False, + True, +}; + +inline bool IsRetryEnabled(const TRetryOperationSettings& settings) { + return settings.MaxRetries_ > 0; +} + +inline TRetryOperationSettings ResolveRetrySettings( + const TRetryOperationSettings& clientDefault, + const std::optional& operationOverride, + const std::optional& explicitOverride, + TDuration operationClientTimeout, + ERetryIdempotentDefault idempotentDefault) +{ + TRetryOperationSettings settings = explicitOverride.value_or( + operationOverride.value_or(clientDefault)); + + if (operationClientTimeout != TDuration::Max()) { + if (settings.MaxTimeout_ == TDuration::Max() || operationClientTimeout < settings.MaxTimeout_) { + settings.MaxTimeout(operationClientTimeout); + } + } + + if (!explicitOverride && !operationOverride && idempotentDefault == ERetryIdempotentDefault::True + && !clientDefault.IdempotentWasSet_) + { + settings.Idempotent(true); + } + + return settings; +} + +inline TRetryOperationSettings ResolveRetrySettings( + const TRetryOperationSettings& clientDefault, + const std::optional& operationOverride, + TDuration operationClientTimeout, + ERetryIdempotentDefault idempotentDefault) +{ + return ResolveRetrySettings( + clientDefault, operationOverride, std::nullopt, operationClientTimeout, idempotentDefault); +} + +template +auto RunUnaryWithRetry(TClient& client, TRetryOperationSettings settings, TRunOnce&& runOnce) + -> decltype(runOnce(TDuration::Max())) +{ + if (client.GetInRetryOperationContext()) { + return runOnce(TDuration::Max()); + } + if (!IsRetryEnabled(settings)) { + return runOnce(settings.MaxTimeout_); + } + + using TResult = decltype(runOnce(TDuration::Max())); + + auto operation = [runOnce = std::forward(runOnce)](TClient& /*clientRef*/, TDuration remainingTimeout) -> TResult { + return runOnce(remainingTimeout); + }; + + using TRetryAsync = Async::TRetryWithoutSession; + using TRetryContextAsync = Async::TRetryContext; + + return typename TRetryContextAsync::TPtr(new TRetryAsync(client, std::move(operation), settings))->Execute(); +} + +} // namespace NYdb::NRetry diff --git a/src/client/impl/internal/retry/retry_sync.h b/src/client/impl/internal/retry/retry_sync.h index 2fdf84bf6d8..e21b3c5eae6 100644 --- a/src/client/impl/internal/retry/retry_sync.h +++ b/src/client/impl/internal/retry/retry_sync.h @@ -130,11 +130,14 @@ class TRetryWithoutSession : public TRetryContext { } TStatusType RunOperation() override { - if constexpr (TFunctionArgs::Length == 1) { - return Operation_(this->Client_); - } else { - return Operation_(this->Client_, this->GetRemainingTimeout()); - } + return InvokeWithRangeErrorCatch([&]() -> TStatusType { + TInRetryOperationContextClientGuard guard(this->Client_); + if constexpr (TFunctionArgs::Length == 1) { + return Operation_(this->Client_); + } else { + return Operation_(this->Client_, this->GetRemainingTimeout()); + } + }); } }; @@ -169,7 +172,7 @@ class TRetryWithSession : public TRetryContext, public TRe Session_ = sessionResult.GetSession(); TRetryDeadlineHelper::SetDeadline(*Session_, Deadline_); } - status = TStatusType(TStatus(sessionResult)); + status = MakeRetryResultFromStatus(TStatus(sessionResult)); } if (Session_) { @@ -180,11 +183,14 @@ class TRetryWithSession : public TRetryContext, public TRe } TStatusType RunOperation() override { - if constexpr (TFunctionArgs::Length == 1) { - return Operation_(this->Session_.value()); - } else { - return Operation_(this->Session_.value(), this->GetRemainingTimeout()); - } + return InvokeWithRangeErrorCatch([&]() -> TStatusType { + TInRetryOperationContextClientGuard guard(this->Client_); + if constexpr (TFunctionArgs::Length == 1) { + return Operation_(this->Session_.value()); + } else { + return Operation_(this->Session_.value(), this->GetRemainingTimeout()); + } + }); } void Reset() override { diff --git a/src/client/impl/observability/CMakeLists.txt b/src/client/impl/observability/CMakeLists.txt index bca5e364657..5c69ccc3301 100644 --- a/src/client/impl/observability/CMakeLists.txt +++ b/src/client/impl/observability/CMakeLists.txt @@ -13,6 +13,7 @@ target_link_libraries(impl-observability PUBLIC target_sources(impl-observability PRIVATE metrics.cpp + metric_buffer.cpp observation.cpp span.cpp ) diff --git a/src/client/impl/session/session_client.h b/src/client/impl/session/session_client.h index bd0484f7d8b..5054a690b46 100644 --- a/src/client/impl/session/session_client.h +++ b/src/client/impl/session/session_client.h @@ -1,5 +1,7 @@ #pragma once +#include + namespace NYdb::inline V3 { class TKqpSessionCommon; @@ -10,6 +12,8 @@ class ISessionClient { virtual void DeleteSession(TKqpSessionCommon* sessionImpl) = 0; + virtual void PessimizeNode(std::uint64_t nodeId) = 0; + // TODO: Try to remove from ISessionClient virtual bool ReturnSession(TKqpSessionCommon* sessionImpl) = 0; }; diff --git a/src/client/operation/operation.cpp b/src/client/operation/operation.cpp index aaeb5791fe6..e1354864242 100644 --- a/src/client/operation/operation.cpp +++ b/src/client/operation/operation.cpp @@ -91,4 +91,10 @@ NThreading::TFuture> TOperationCli return List("compaction", pageSize, pageToken); } +template NThreading::TFuture TOperationClient::Get(const TOperation::TOperationId& id); +template <> +NThreading::TFuture> TOperationClient::List(std::uint64_t pageSize, const std::string& pageToken) { + return List("analyze", pageSize, pageToken); +} + } // namespace NYdb::NOperation diff --git a/src/client/persqueue_public/impl/persqueue.cpp b/src/client/persqueue_public/impl/persqueue.cpp index 4829a76a6a6..0188c5d749e 100644 --- a/src/client/persqueue_public/impl/persqueue.cpp +++ b/src/client/persqueue_public/impl/persqueue.cpp @@ -10,6 +10,45 @@ namespace NYdb::inline V3::NPersQueue { +namespace { + +TDuration ProtoDurationToTDuration(const google::protobuf::Duration& duration) { + return TDuration::Seconds(duration.seconds()) + TDuration::MicroSeconds(duration.nanos() / 1000); +} + +TSharedConsumerSettings ConvertProtoToSharedConsumerSettings( + const Ydb::PersQueue::V1::TopicSettings::SharedConsumerType& shared) +{ + TSharedConsumerSettings settings; + settings.KeepMessagesOrder(shared.keep_messages_order()); + if (shared.has_default_processing_timeout()) { + settings.DefaultProcessingTimeout(ProtoDurationToTDuration(shared.default_processing_timeout())); + } + if (shared.has_receive_message_wait_time()) { + settings.ReceiveMessageWaitTime(ProtoDurationToTDuration(shared.receive_message_wait_time())); + } + if (shared.has_receive_message_delay()) { + settings.ReceiveMessageDelay(ProtoDurationToTDuration(shared.receive_message_delay())); + } + if (shared.has_dead_letter_policy()) { + const auto& policy = shared.dead_letter_policy(); + TSharedConsumerDeadLetterPolicySettings deadLetterPolicy; + deadLetterPolicy.Enabled(policy.enabled()); + if (policy.has_condition()) { + deadLetterPolicy.MaxProcessingAttempts(policy.condition().max_processing_attempts()); + } + if (policy.has_delete_action()) { + deadLetterPolicy.DeleteAction(); + } else if (policy.has_move_action()) { + deadLetterPolicy.MoveAction(policy.move_action().dead_letter_queue()); + } + settings.DeadLetterPolicy(deadLetterPolicy); + } + return settings; +} + +} // namespace + class TCommonCodecsProvider { public: TCommonCodecsProvider() { @@ -150,6 +189,9 @@ TDescribeTopicResult::TTopicSettings::TReadRule::TReadRule(const Ydb::PersQueue: } Version_ = settings.version(); ServiceType_ = settings.service_type(); + if (settings.has_shared_consumer_type()) { + SharedConsumer_ = ConvertProtoToSharedConsumerSettings(settings.shared_consumer_type()); + } } TDescribeTopicResult::TTopicSettings::TRemoteMirrorRule::TRemoteMirrorRule(const Ydb::PersQueue::V1::TopicSettings::RemoteMirrorRule& settings) diff --git a/src/client/persqueue_public/impl/persqueue_impl.h b/src/client/persqueue_public/impl/persqueue_impl.h index 01baccc86f9..18ecfa2cdf0 100644 --- a/src/client/persqueue_public/impl/persqueue_impl.h +++ b/src/client/persqueue_public/impl/persqueue_impl.h @@ -33,13 +33,53 @@ class TPersQueueClient::TImpl : public TClientImplCommonset_seconds(duration.Seconds()); + proto->set_nanos(duration.NanoSecondsOfSecond()); + } + + static void ConvertSharedConsumerToProto( + const TSharedConsumerSettings& settings, + Ydb::PersQueue::V1::TopicSettings::SharedConsumerType& shared) + { + shared.set_keep_messages_order(settings.KeepMessagesOrder_); + if (settings.DefaultProcessingTimeout_ != TDuration::Zero()) { + SetProtoDuration(shared.mutable_default_processing_timeout(), settings.DefaultProcessingTimeout_); + } + if (settings.ReceiveMessageWaitTime_ != TDuration::Zero()) { + SetProtoDuration(shared.mutable_receive_message_wait_time(), settings.ReceiveMessageWaitTime_); + } + if (settings.ReceiveMessageDelay_ != TDuration::Zero()) { + SetProtoDuration(shared.mutable_receive_message_delay(), settings.ReceiveMessageDelay_); + } + + const auto& deadLetterPolicy = settings.DeadLetterPolicy_; + if (deadLetterPolicy.Enabled_) { + auto& protoDeadLetterPolicy = *shared.mutable_dead_letter_policy(); + protoDeadLetterPolicy.set_enabled(true); + if (deadLetterPolicy.MaxProcessingAttempts_) { + protoDeadLetterPolicy.mutable_condition()->set_max_processing_attempts(deadLetterPolicy.MaxProcessingAttempts_); + } + switch (deadLetterPolicy.Action_) { + case TSharedConsumerDeadLetterPolicySettings::EAction::Move: + protoDeadLetterPolicy.mutable_move_action()->set_dead_letter_queue( + TStringType{deadLetterPolicy.DeadLetterQueue_}); + break; + case TSharedConsumerDeadLetterPolicySettings::EAction::Delete: + protoDeadLetterPolicy.mutable_delete_action(); + break; + case TSharedConsumerDeadLetterPolicySettings::EAction::Unspecified: + break; + } + } + } + template static void ConvertToProtoReadRule(const TReadRule& readRule, Ydb::PersQueue::V1::TopicSettings::ReadRule& rrProps) { rrProps.set_consumer_name(TStringType{readRule.ConsumerName_}); rrProps.set_important(readRule.Important_); if (readRule.AvailabilityPeriod_ != TDuration::Zero()) { - rrProps.mutable_availability_period()->set_seconds(readRule.AvailabilityPeriod_.Seconds()); - rrProps.mutable_availability_period()->set_nanos(readRule.AvailabilityPeriod_.NanoSecondsOfSecond()); + SetProtoDuration(rrProps.mutable_availability_period(), readRule.AvailabilityPeriod_); } rrProps.set_starting_message_timestamp_ms(readRule.StartingMessageTimestamp_.MilliSeconds()); rrProps.set_version(readRule.Version_); @@ -48,6 +88,9 @@ class TPersQueueClient::TImpl : public TClientImplCommon(codec))); } rrProps.set_service_type(TStringType{readRule.ServiceType_}); + if (readRule.SharedConsumer_) { + ConvertSharedConsumerToProto(readRule.SharedConsumer_.value(), *rrProps.mutable_shared_consumer_type()); + } } template diff --git a/src/client/persqueue_public/impl/read_session.cpp b/src/client/persqueue_public/impl/read_session.cpp index bbb8f06e996..9323a353e12 100644 --- a/src/client/persqueue_public/impl/read_session.cpp +++ b/src/client/persqueue_public/impl/read_session.cpp @@ -557,7 +557,7 @@ TReadSessionEvent::TCreatePartitionStreamEvent::TCreatePartitionStreamEvent(TPar void TReadSessionEvent::TCreatePartitionStreamEvent::Confirm(std::optional readOffset, std::optional commitOffset) { if (PartitionStream) { - static_cast(PartitionStream.Get())->ConfirmCreate(readOffset, commitOffset); + static_cast(PartitionStream.Get())->ConfirmCreate(readOffset, commitOffset, {}); } } diff --git a/src/client/persqueue_public/impl/write_session_impl.cpp b/src/client/persqueue_public/impl/write_session_impl.cpp index 44915db0586..2029893d34e 100644 --- a/src/client/persqueue_public/impl/write_session_impl.cpp +++ b/src/client/persqueue_public/impl/write_session_impl.cpp @@ -13,6 +13,8 @@ namespace NYdb::inline V3::NPersQueue { +namespace NGrpc = NTopic::NWriteSessionGrpc; + const TDuration UPDATE_TOKEN_PERIOD = TDuration::Hours(1); //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -30,8 +32,8 @@ TWriteSessionImpl::TWriteSessionImpl( , PrevToken(DbDriverState->CredentialsProvider ? DbDriverState->CredentialsProvider->GetAuthInfo() : "") , InitSeqNoPromise(NThreading::NewPromise()) , WakeupInterval( - Settings.BatchFlushInterval_.value_or(TDuration::Zero()) ? - std::min(Settings.BatchFlushInterval_.value_or(TDuration::Seconds(1)) / 5, TDuration::MilliSeconds(100)) + Settings.BatchFlushInterval_ != TDuration::Zero() ? + std::min(Settings.BatchFlushInterval_ / 5, TDuration::MilliSeconds(100)) : TDuration::MilliSeconds(100) ) @@ -992,7 +994,7 @@ void TWriteSessionImpl::FlushWriteIfRequiredImpl() { if (!CurrentBatch.Empty() && !CurrentBatch.FlushRequested) { MessagesAcquired += static_cast(CurrentBatch.Acquire()); - if (TInstant::Now() - CurrentBatch.StartedAt >= Settings.BatchFlushInterval_.value_or(TDuration::Zero()) + if (TInstant::Now() - CurrentBatch.StartedAt >= Settings.BatchFlushInterval_ || CurrentBatch.CurrentSize >= Settings.BatchFlushSizeBytes_.value_or(0) || CurrentBatch.CurrentSize >= MaxBlockSize || CurrentBatch.Messages.size() >= MaxBlockMessageCount @@ -1062,8 +1064,33 @@ size_t TWriteSessionImpl::WriteBatchImpl() { return size; } -size_t GetMaxGrpcMessageSize() { - return 120_MB; +size_t EstimatePackedInt64FieldUpperBound(ui32 fieldNumber, size_t itemsCount) { + static constexpr size_t MaxInt64VarintPayloadSize = 10; + + return NGrpc::ProtoPackedInt64FieldSize(fieldNumber, itemsCount * MaxInt64VarintPayloadSize); +} + +template +size_t EstimateWriteRequestBlockSize(const TBlock& block) { + size_t size = 0; + size += EstimatePackedInt64FieldUpperBound(2, block.MessageCount); // sequence_numbers + size += EstimatePackedInt64FieldUpperBound(3, block.MessageCount); // created_at_ms + size += EstimatePackedInt64FieldUpperBound(4, block.MessageCount); // sent_at_ms + size += EstimatePackedInt64FieldUpperBound(5, block.MessageCount); // message_sizes + + size += NGrpc::ProtoPackedInt64FieldSize(6, NGrpc::TWireFormatLite::Int64Size(static_cast(block.Offset))); + size += NGrpc::ProtoPackedInt64FieldSize(7, NGrpc::TWireFormatLite::Int64Size(static_cast(block.PartNumber))); + size += NGrpc::ProtoPackedInt64FieldSize(8, NGrpc::TWireFormatLite::Int64Size(static_cast(block.MessageCount))); + size += NGrpc::ProtoPackedInt64FieldSize(9, NGrpc::TWireFormatLite::Int64Size(static_cast(block.OriginalSize))); + size += NGrpc::ProtoBytesFieldSize(10, block.CodecID.size()); + if (block.Compressed) { + size += NGrpc::ProtoBytesFieldSize(11, block.Data.size()); + } else { + for (const auto& buffer : block.OriginalDataRefs) { + size += NGrpc::ProtoBytesFieldSize(11, buffer.size()); + } + } + return size; } bool TWriteSessionImpl::IsReadyToSendNextImpl() { @@ -1180,11 +1207,18 @@ void TWriteSessionImpl::SendImpl() { TClientMessage clientMessage; auto* writeRequest = clientMessage.mutable_write_request(); auto sentAtMs = TInstant::Now().MilliSeconds(); + NGrpc::TRequestSizeLimiter sizeLimiter(2); // Sent blocks while we can without messages reordering - while (IsReadyToSendNextImpl() && clientMessage.ByteSizeLong() < GetMaxGrpcMessageSize()) { + while (IsReadyToSendNextImpl()) { const auto& block = PackedMessagesToSend.top(); Y_ABORT_UNLESS(block.Valid); + + const size_t blockSize = EstimateWriteRequestBlockSize(block); + if (!sizeLimiter.CanAdd(blockSize)) { + break; + } + for (size_t i = 0; i != block.MessageCount; ++i) { Y_ABORT_UNLESS(!OriginalMessagesToSend.empty()); @@ -1216,6 +1250,7 @@ void TWriteSessionImpl::SendImpl() { moveBlock.Move(block); SentPackedMessage.emplace(std::move(moveBlock)); PackedMessagesToSend.pop(); + sizeLimiter.Add(blockSize); } UpdateTokenIfNeededImpl(); LOG_LAZY(DbDriverState->Log, diff --git a/src/client/persqueue_public/include/control_plane.h b/src/client/persqueue_public/include/control_plane.h index 3d1413b00a0..19d865c78f1 100644 --- a/src/client/persqueue_public/include/control_plane.h +++ b/src/client/persqueue_public/include/control_plane.h @@ -42,6 +42,53 @@ struct TCredentials { Ydb::PersQueue::V1::Credentials Credentials_; }; +struct TSharedConsumerDeadLetterPolicySettings { + using TSelf = TSharedConsumerDeadLetterPolicySettings; + + enum class EAction { + Unspecified = 0, + Move = 1, + Delete = 2, + }; + + FLUENT_SETTING_DEFAULT(bool, Enabled, false); + FLUENT_SETTING_DEFAULT(ui32, MaxProcessingAttempts, 0); + FLUENT_SETTING(std::string, DeadLetterQueue); + FLUENT_SETTING_DEFAULT(EAction, Action, EAction::Unspecified); + + bool GetEnabled() const { return Enabled_; } + ui32 GetMaxProcessingAttempts() const { return MaxProcessingAttempts_; } + const std::string& GetDeadLetterQueue() const { return DeadLetterQueue_; } + EAction GetAction() const { return Action_; } + + TSelf& DeleteAction() { + Action_ = EAction::Delete; + return *this; + } + + TSelf& MoveAction(const std::string& deadLetterQueue) { + Action_ = EAction::Move; + DeadLetterQueue_ = deadLetterQueue; + return *this; + } +}; + +struct TSharedConsumerSettings { + using TSelf = TSharedConsumerSettings; + + FLUENT_SETTING_DEFAULT(bool, KeepMessagesOrder, false); + FLUENT_SETTING_DEFAULT(TDuration, DefaultProcessingTimeout, TDuration::Zero()); + FLUENT_SETTING_DEFAULT(TDuration, ReceiveMessageWaitTime, TDuration::Zero()); + FLUENT_SETTING_DEFAULT(TDuration, ReceiveMessageDelay, TDuration::Zero()); + FLUENT_SETTING(TSharedConsumerDeadLetterPolicySettings, DeadLetterPolicy); + + bool GetKeepMessagesOrder() const { return KeepMessagesOrder_; } + TDuration GetDefaultProcessingTimeout() const { return DefaultProcessingTimeout_; } + TDuration GetReceiveMessageWaitTime() const { return ReceiveMessageWaitTime_; } + TDuration GetReceiveMessageDelay() const { return ReceiveMessageDelay_; } + const TSharedConsumerDeadLetterPolicySettings& GetDeadLetterPolicy() const { return DeadLetterPolicy_; } +}; + // Result for describe resource request. struct TDescribeTopicResult : public TStatus { @@ -65,6 +112,7 @@ struct TDescribeTopicResult : public TStatus { } GETTER(ui32, Version); GETTER(std::string, ServiceType); + const std::optional& SharedConsumer() const { return SharedConsumer_; } private: std::string ConsumerName_; @@ -75,6 +123,7 @@ struct TDescribeTopicResult : public TStatus { std::vector SupportedCodecs_; ui32 Version_; std::string ServiceType_; + std::optional SharedConsumer_; }; struct TRemoteMirrorRule { @@ -190,6 +239,9 @@ struct TReadRuleSettings { FLUENT_SETTING_DEFAULT(ui32, Version, 0); FLUENT_SETTING(std::string, ServiceType); + FLUENT_SETTING_OPTIONAL(TSharedConsumerSettings, SharedConsumer); + + const std::optional& GetSharedConsumer() const { return SharedConsumer_; } TReadRuleSettings& SetSettings(const TDescribeTopicResult::TTopicSettings::TReadRule& settings) { ConsumerName_ = settings.ConsumerName(); @@ -203,6 +255,10 @@ struct TReadRuleSettings { } Version_ = settings.Version(); ServiceType_ = settings.ServiceType(); + SharedConsumer_.reset(); + if (const auto& sharedConsumer = settings.SharedConsumer()) { + SharedConsumer_ = sharedConsumer; + } return *this; } diff --git a/src/client/persqueue_public/include/write_session.h b/src/client/persqueue_public/include/write_session.h index 5f7c0d49d00..29a3ec69ee5 100644 --- a/src/client/persqueue_public/include/write_session.h +++ b/src/client/persqueue_public/include/write_session.h @@ -77,9 +77,8 @@ struct TWriteSessionSettings : public TRequestSettings { //! but for no longer than BatchFlushInterval. //! Upon reaching FlushInterval or FlushSize limit, all messages will be written with one batch. //! Greatly increases performance for small messages. - //! Setting either value to zero means immediate write with no batching. (Unrecommended, especially for clients - //! sending small messages at high rate). - FLUENT_SETTING_OPTIONAL(TDuration, BatchFlushInterval); + //! Zero BatchFlushInterval or BatchFlushSizeBytes disables the corresponding limit (immediate flush). + FLUENT_SETTING_DEFAULT(TDuration, BatchFlushInterval, TDuration::Seconds(1)); FLUENT_SETTING_OPTIONAL(ui64, BatchFlushSizeBytes); FLUENT_SETTING_DEFAULT(TDuration, ConnectTimeout, TDuration::Seconds(30)); diff --git a/src/client/persqueue_public/ut/read_session_ut.cpp b/src/client/persqueue_public/ut/read_session_ut.cpp index 63f60b3909b..92b45da0a53 100644 --- a/src/client/persqueue_public/ut/read_session_ut.cpp +++ b/src/client/persqueue_public/ut/read_session_ut.cpp @@ -1862,7 +1862,8 @@ Y_UNIT_TEST_SUITE(ReadSessionImplTest) { \ auto event = sessionQueue.GetEventImpl(maxByteSize, accumulator); \ \ - UNIT_ASSERT(std::holds_alternative(event.GetEvent()));\ + UNIT_ASSERT(event);\ + UNIT_ASSERT(std::holds_alternative(event->GetEvent()));\ } #define UNIT_ASSERT_DATA_EVENT(count) \ @@ -1874,8 +1875,9 @@ Y_UNIT_TEST_SUITE(ReadSessionImplTest) { \ auto event = sessionQueue.GetEventImpl(maxByteSize, accumulator); \ \ - UNIT_ASSERT(std::holds_alternative(event.GetEvent())); \ - UNIT_ASSERT_VALUES_EQUAL(std::get(event.GetEvent()).GetMessagesCount(), count); \ + UNIT_ASSERT(event);\ + UNIT_ASSERT(std::holds_alternative(event->GetEvent())); \ + UNIT_ASSERT_VALUES_EQUAL(std::get(event->GetEvent()).GetMessagesCount(), count); \ } NTopic::TAReadSessionSettings settings; diff --git a/src/client/query/CMakeLists.txt b/src/client/query/CMakeLists.txt index 3cc7401200b..d0f1b746aba 100644 --- a/src/client/query/CMakeLists.txt +++ b/src/client/query/CMakeLists.txt @@ -11,6 +11,7 @@ target_link_libraries(client-ydb_query PUBLIC client-ydb_driver client-ydb_query-impl client-ydb_result + client-ydb_row_ranges client-metrics client-trace client-types-operation diff --git a/src/client/query/client.cpp b/src/client/query/client.cpp index 2260d9f401d..4ff31745aa0 100644 --- a/src/client/query/client.cpp +++ b/src/client/query/client.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -58,6 +59,9 @@ static void SetTxSettings(const TTxSettings& txSettings, Ydb::Query::Transaction case TTxSettings::TS_READ_COMMITTED_RW: proto->mutable_read_committed_read_write(); break; + case TTxSettings::TS_STRICT_SERIALIZABLE_RW: + proto->mutable_strict_serializable_read_write(); + break; default: throw TContractViolation("Unexpected transaction mode."); } @@ -374,9 +378,6 @@ class TQueryClient::TImpl: public TClientImplCommon, public } void DeleteSession(TKqpSessionCommon* sessionImpl) override { - //TODO: Remove this copy-paste - - // Closing not owned by session pool session should not fire getting new session if (sessionImpl->IsOwnedBySessionPool()) { if (SessionPool_.CheckAndFeedWaiterNewSession(sessionImpl->NeedUpdateActiveCounter())) { // We requested new session for waiter which already incremented @@ -408,6 +409,10 @@ class TQueryClient::TImpl: public TClientImplCommon, public return true; } + void PessimizeNode(std::uint64_t nodeId) override { + DbDriverState_->EndpointPool.BanNodeId(nodeId); + } + void DoAttachSession(Ydb::Query::CreateSessionResponse* resp , NThreading::TPromise promise , const std::string& endpoint @@ -645,6 +650,8 @@ class TQueryClient::TImpl: public TClientImplCommon, public ); } + TClientSettings Settings_; + private: std::shared_ptr MakeObservation(const std::string& operationName) { return std::make_shared( @@ -662,7 +669,6 @@ class TQueryClient::TImpl: public TClientImplCommon, public NSdkStats::TAtomicHistogram<::NMonitoring::THistogram> QuerySizeHistogram_; NSdkStats::TAtomicHistogram<::NMonitoring::THistogram> ParamsSizeHistogram_; - TClientSettings Settings_; NSessionPool::TSessionPool SessionPool_; }; @@ -675,13 +681,39 @@ TQueryClient::TQueryClient(const TDriver& driver, const TClientSettings& setting TAsyncExecuteQueryResult TQueryClient::ExecuteQuery(const std::string& query, const TTxControl& txControl, const TExecuteQuerySettings& settings) { - return Impl_->ExecuteQuery(query, txControl, {}, settings); + const auto retrySettings = NRetry::ResolveRetrySettings( + Impl_->Settings_.RetrySettings_, + settings.RetrySettings_, + settings.ClientTimeout_, + NRetry::ERetryIdempotentDefault::False); + + return NRetry::RunUnaryWithRetry(*this, retrySettings, + [this, query, txControl, settings](TDuration timeout) { + auto opSettings = settings; + if (timeout != TDuration::Max()) { + opSettings.ClientTimeout(timeout); + } + return Impl_->ExecuteQuery(query, txControl, {}, opSettings); + }); } TAsyncExecuteQueryResult TQueryClient::ExecuteQuery(const std::string& query, const TTxControl& txControl, const TParams& params, const TExecuteQuerySettings& settings) { - return Impl_->ExecuteQuery(query, txControl, params, settings); + const auto retrySettings = NRetry::ResolveRetrySettings( + Impl_->Settings_.RetrySettings_, + settings.RetrySettings_, + settings.ClientTimeout_, + NRetry::ERetryIdempotentDefault::False); + + return NRetry::RunUnaryWithRetry(*this, retrySettings, + [this, query, txControl, params, settings](TDuration timeout) { + auto opSettings = settings; + if (timeout != TDuration::Max()) { + opSettings.ClientTimeout(timeout); + } + return Impl_->ExecuteQuery(query, txControl, params, opSettings); + }); } TAsyncExecuteQueryIterator TQueryClient::StreamExecuteQuery(const std::string& query, const TTxControl& txControl, @@ -697,21 +729,66 @@ TAsyncExecuteQueryIterator TQueryClient::StreamExecuteQuery(const std::string& q } NThreading::TFuture TQueryClient::ExecuteScript(const std::string& script, - const TExecuteScriptSettings& settings) + const TExecuteScriptSettings& settings, + const std::optional& retrySettings) { - return Impl_->ExecuteScript(script, {}, settings); + const auto resolvedRetrySettings = NRetry::ResolveRetrySettings( + Impl_->Settings_.RetrySettings_, + settings.RetrySettings_, + retrySettings, + settings.ClientTimeout_, + NRetry::ERetryIdempotentDefault::False); + + return NRetry::RunUnaryWithRetry(*this, resolvedRetrySettings, + [this, script, settings](TDuration timeout) { + auto opSettings = settings; + if (timeout != TDuration::Max()) { + opSettings.ClientTimeout(timeout); + } + return Impl_->ExecuteScript(script, {}, opSettings); + }); } NThreading::TFuture TQueryClient::ExecuteScript(const std::string& script, - const TParams& params, const TExecuteScriptSettings& settings) + const TParams& params, const TExecuteScriptSettings& settings, + const std::optional& retrySettings) { - return Impl_->ExecuteScript(script, params, settings); + const auto resolvedRetrySettings = NRetry::ResolveRetrySettings( + Impl_->Settings_.RetrySettings_, + settings.RetrySettings_, + retrySettings, + settings.ClientTimeout_, + NRetry::ERetryIdempotentDefault::False); + + return NRetry::RunUnaryWithRetry(*this, resolvedRetrySettings, + [this, script, params, settings](TDuration timeout) { + auto opSettings = settings; + if (timeout != TDuration::Max()) { + opSettings.ClientTimeout(timeout); + } + return Impl_->ExecuteScript(script, params, opSettings); + }); } TAsyncFetchScriptResultsResult TQueryClient::FetchScriptResults(const NKikimr::NOperationId::TOperationId& operationId, int64_t resultSetIndex, - const TFetchScriptResultsSettings& settings) + const TFetchScriptResultsSettings& settings, + const std::optional& retrySettings) { - return Impl_->FetchScriptResults(operationId, resultSetIndex, settings); + const auto resolvedRetrySettings = NRetry::ResolveRetrySettings( + Impl_->Settings_.RetrySettings_, + settings.RetrySettings_, + retrySettings, + settings.ClientTimeout_, + NRetry::ERetryIdempotentDefault::True); + + return NRetry::RunUnaryWithRetry(*this, resolvedRetrySettings, + [this, operationId, resultSetIndex, settings](TDuration timeout) { + auto opSettings = settings; + if (timeout != TDuration::Max()) { + opSettings.ClientTimeout(timeout); + } + return Impl_->FetchScriptResults(operationId, resultSetIndex, opSettings); + }); } TAsyncCreateSessionResult TQueryClient::GetSession(const TCreateSessionSettings& settings) @@ -731,6 +808,18 @@ int64_t TQueryClient::GetCurrentPoolSize() const { return Impl_->GetCurrentPoolSize(); } +namespace { +thread_local bool QueryClientInRetryOperationContext = false; +} // namespace + +bool TQueryClient::GetInRetryOperationContext() const { + return QueryClientInRetryOperationContext; +} + +void TQueryClient::SetInRetryOperationContext(bool value) { + QueryClientInRetryOperationContext = value; +} + TAsyncExecuteQueryResult TQueryClient::RetryQuery(TQueryResultFunc&& queryFunc, TRetryOperationSettings settings) { return TRetryContextResultAsync::TPtr( diff --git a/src/client/query/impl/CMakeLists.txt b/src/client/query/impl/CMakeLists.txt index d33258b7afd..1715df4866b 100644 --- a/src/client/query/impl/CMakeLists.txt +++ b/src/client/query/impl/CMakeLists.txt @@ -16,6 +16,7 @@ target_link_libraries(client-ydb_query-impl PUBLIC target_sources(client-ydb_query-impl PRIVATE exec_query.cpp client_session.cpp + session_state_handler.cpp ) _ydb_sdk_install_targets(TARGETS client-ydb_query-impl) diff --git a/src/client/query/impl/client_session.cpp b/src/client/query/impl/client_session.cpp index 7492e0d13e3..6049ecb1ac6 100644 --- a/src/client/query/impl/client_session.cpp +++ b/src/client/query/impl/client_session.cpp @@ -1,4 +1,5 @@ #include "client_session.h" +#include "session_state_handler.h" #define INCLUDE_YDB_INTERNAL_H #include @@ -42,7 +43,7 @@ class TSafeTSessionImplHolder { {} TSession::TImpl* TrySharedOwning() noexcept { - auto old = Semaphore.fetch_add(1); + auto old = Semaphore.fetch_add(1); if (old == 0) { OwnerThread.store(std::this_thread::get_id()); return Ptr; @@ -76,9 +77,17 @@ void TSession::TImpl::StartAsyncRead(TStreamProcessorPtr ptr, std::weak_ptr(); ptr->Read(resp.get(), [resp, ptr, client, holder](NYdbGrpc::TGrpcStatus grpcStatus) mutable { switch (grpcStatus.GRpcStatusCode) { - case grpc::StatusCode::OK: - StartAsyncRead(ptr, client, holder); + case grpc::StatusCode::OK: { + auto impl = holder->TrySharedOwning(); + if (impl) { + const auto action = HandleAttachSessionState(*resp, impl, client.lock()); + if (action == EAttachStreamReadAction::Continue) { + StartAsyncRead(ptr, client, holder); + } + holder->Release(); + } break; + } default: { auto impl = holder->TrySharedOwning(); if (impl) { diff --git a/src/client/query/impl/exec_query.cpp b/src/client/query/impl/exec_query.cpp index 0301ea5968e..0a7814b73d4 100644 --- a/src/client/query/impl/exec_query.cpp +++ b/src/client/query/impl/exec_query.cpp @@ -40,6 +40,9 @@ static void SetTxSettings(const TTxSettings& txSettings, Ydb::Query::Transaction case TTxSettings::TS_READ_COMMITTED_RW: proto->mutable_read_committed_read_write(); break; + case TTxSettings::TS_STRICT_SERIALIZABLE_RW: + proto->mutable_strict_serializable_read_write(); + break; default: throw TContractViolation("Unexpected transaction mode."); } diff --git a/src/client/query/impl/session_state_handler.cpp b/src/client/query/impl/session_state_handler.cpp new file mode 100644 index 00000000000..a5e64063b17 --- /dev/null +++ b/src/client/query/impl/session_state_handler.cpp @@ -0,0 +1,33 @@ +#include "session_state_handler.h" + +namespace NYdb::inline V3::NQuery { + +EAttachStreamReadAction HandleAttachSessionState( + const Ydb::Query::SessionState& state, + TKqpSessionCommon* session, + const std::shared_ptr& client) +{ + if (state.has_session_shutdown() || state.has_node_shutdown()) { + if (!session) { + return EAttachStreamReadAction::Stop; + } + if (state.has_node_shutdown()) { + const auto nodeId = session->GetEndpointKey().GetNodeId(); + if (nodeId != 0 && client) { + client->PessimizeNode(nodeId); + } + } + if (session->GetState() == TKqpSessionCommon::S_IDLE) { + if (client) { + session->CloseFromServer(client); + } + } else { + session->MarkAsClosing(); + } + return EAttachStreamReadAction::Stop; + } + + return EAttachStreamReadAction::Continue; +} + +} diff --git a/src/client/query/impl/session_state_handler.h b/src/client/query/impl/session_state_handler.h new file mode 100644 index 00000000000..9934d122416 --- /dev/null +++ b/src/client/query/impl/session_state_handler.h @@ -0,0 +1,21 @@ +#pragma once + +#include + +#include + +#include + +namespace NYdb::inline V3::NQuery { + +enum class EAttachStreamReadAction { + Continue, + Stop, +}; + +EAttachStreamReadAction HandleAttachSessionState( + const Ydb::Query::SessionState& state, + TKqpSessionCommon* session, + const std::shared_ptr& client); + +} diff --git a/src/client/query/stats.cpp b/src/client/query/stats.cpp index fc82e4293d0..22a47b21ebe 100644 --- a/src/client/query/stats.cpp +++ b/src/client/query/stats.cpp @@ -15,6 +15,113 @@ class TExecStats::TImpl { Ydb::TableStats::QueryStats Proto; }; +TOperationStats::TOperationStats(const Ydb::TableStats::OperationStats& proto) + : Rows_(proto.rows()) + , Bytes_(proto.bytes()) +{} + +uint64_t TOperationStats::GetRows() const { + return Rows_; +} + +uint64_t TOperationStats::GetBytes() const { + return Bytes_; +} + +TTableAccessStats::TTableAccessStats(const Ydb::TableStats::TableAccessStats& proto) + : Name_(proto.name()) + , Reads_(proto.reads()) + , Updates_(proto.updates()) + , Deletes_(proto.deletes()) + , PartitionsCount_(proto.partitions_count()) +{} + +const std::string& TTableAccessStats::GetName() const { + return Name_; +} + +const TOperationStats& TTableAccessStats::GetReads() const { + return Reads_; +} + +const TOperationStats& TTableAccessStats::GetUpdates() const { + return Updates_; +} + +const TOperationStats& TTableAccessStats::GetDeletes() const { + return Deletes_; +} + +uint64_t TTableAccessStats::GetPartitionsCount() const { + return PartitionsCount_; +} + +TQueryPhaseStats::TQueryPhaseStats(const Ydb::TableStats::QueryPhaseStats& proto) + : DurationUs_(proto.duration_us()) + , CpuTimeUs_(proto.cpu_time_us()) + , AffectedShards_(proto.affected_shards()) + , LiteralPhase_(proto.literal_phase()) +{ + TableAccess_.reserve(proto.table_access().size()); + for (const auto& tableAccess : proto.table_access()) { + TableAccess_.emplace_back(tableAccess); + } +} + +uint64_t TQueryPhaseStats::GetDurationUs() const { + return DurationUs_; +} + +TDuration TQueryPhaseStats::GetDuration() const { + return TDuration::MicroSeconds(DurationUs_); +} + +uint64_t TQueryPhaseStats::GetCpuTimeUs() const { + return CpuTimeUs_; +} + +TDuration TQueryPhaseStats::GetCpuTime() const { + return TDuration::MicroSeconds(CpuTimeUs_); +} + +uint64_t TQueryPhaseStats::GetAffectedShards() const { + return AffectedShards_; +} + +bool TQueryPhaseStats::IsLiteralPhase() const { + return LiteralPhase_; +} + +const std::vector& TQueryPhaseStats::GetTableAccess() const { + return TableAccess_; +} + +TCompilationStats::TCompilationStats(const Ydb::TableStats::CompilationStats& proto) + : FromCache_(proto.from_cache()) + , DurationUs_(proto.duration_us()) + , CpuTimeUs_(proto.cpu_time_us()) +{} + +bool TCompilationStats::IsFromCache() const { + return FromCache_; +} + +uint64_t TCompilationStats::GetDurationUs() const { + return DurationUs_; +} + +TDuration TCompilationStats::GetDuration() const { + return TDuration::MicroSeconds(DurationUs_); +} + +uint64_t TCompilationStats::GetCpuTimeUs() const { + return CpuTimeUs_; +} + +TDuration TCompilationStats::GetCpuTime() const { + return TDuration::MicroSeconds(CpuTimeUs_); +} + TExecStats::TExecStats(const Ydb::TableStats::QueryStats& proto) { Impl_ = std::make_shared(); Impl_->Proto = proto; @@ -39,6 +146,18 @@ std::string TExecStats::ToString(bool withPlan) const { return res; } +uint64_t TExecStats::GetProcessCpuTimeUs() const { + return Impl_->Proto.process_cpu_time_us(); +} + +uint64_t TExecStats::GetTotalDurationUs() const { + return Impl_->Proto.total_duration_us(); +} + +uint64_t TExecStats::GetTotalCpuTimeUs() const { + return Impl_->Proto.total_cpu_time_us(); +} + std::optional TExecStats::GetPlan() const { auto proto = Impl_->Proto; @@ -69,6 +188,10 @@ std::optional TExecStats::GetMeta() const { return proto.query_meta(); } +TDuration TExecStats::GetProcessCpuTime() const { + return TDuration::MicroSeconds(Impl_->Proto.process_cpu_time_us()); +} + TDuration TExecStats::GetTotalDuration() const { return TDuration::MicroSeconds(Impl_->Proto.total_duration_us()); } @@ -77,6 +200,23 @@ TDuration TExecStats::GetTotalCpuTime() const { return TDuration::MicroSeconds(Impl_->Proto.total_cpu_time_us()); } +std::vector TExecStats::GetQueryPhases() const { + std::vector phases; + phases.reserve(Impl_->Proto.query_phases().size()); + for (const auto& phase : Impl_->Proto.query_phases()) { + phases.emplace_back(phase); + } + return phases; +} + +std::optional TExecStats::GetCompilation() const { + const auto& proto = Impl_->Proto; + if (!proto.has_compilation()) { + return {}; + } + return TCompilationStats(proto.compilation()); +} + const Ydb::TableStats::QueryStats& TExecStats::GetProto() const { return Impl_->Proto; } diff --git a/src/client/row_ranges/CMakeLists.txt b/src/client/row_ranges/CMakeLists.txt new file mode 100644 index 00000000000..42848e4f0dc --- /dev/null +++ b/src/client/row_ranges/CMakeLists.txt @@ -0,0 +1,15 @@ +_ydb_sdk_add_library(client-ydb_row_ranges) + +target_link_libraries(client-ydb_row_ranges PUBLIC + yutil + threading-future + client-ydb_result + client-ydb_query + client-ydb_table +) + +target_sources(client-ydb_row_ranges PRIVATE + rows.cpp +) + +_ydb_sdk_install_targets(TARGETS client-ydb_row_ranges) diff --git a/src/client/row_ranges/rows.cpp b/src/client/row_ranges/rows.cpp new file mode 100644 index 00000000000..fac682642ea --- /dev/null +++ b/src/client/row_ranges/rows.cpp @@ -0,0 +1,300 @@ +#include + +#include "rows_stream_drain.h" +#include +#include +#include + +#include + +#include +#include +#include + +namespace NYdb::inline V3 { + +namespace { + +template +void ThrowIfUnsuccessful(TStatusCarrier&& carrier) { + if (!carrier.IsSuccess()) { + TStatus status(std::move(carrier)); + throw NStatusHelpers::TYdbRangeErrorException(status) << status; + } +} + +} // namespace + +class TRowRange::TImpl { +public: + struct TEmptyTag {}; + + explicit TImpl(TEmptyTag) + : Producer_(std::make_unique()) + {} + + explicit TImpl(TResultSet&& resultSet) + : Producer_(std::make_unique(std::move(resultSet))) + {} + + explicit TImpl(NQuery::TExecuteQueryIterator&& iterator) + : Producer_(std::make_unique>(std::move(iterator))) + {} + + explicit TImpl(NTable::TScanQueryPartIterator&& iterator) + : Producer_(std::make_unique>(std::move(iterator))) + {} + + explicit TImpl(NTable::TTablePartIterator&& iterator) + : Producer_(std::make_unique>(std::move(iterator))) + {} + + bool Start() { + if (Started_) { + return Parser_.has_value(); + } + Started_ = true; + return AdvanceRow(); + } + + bool Advance() { + Started_ = true; + return AdvanceRow(); + } + + bool HasCurrent() const { + return Parser_.has_value(); + } + + TResultSetParser& CurrentParser() { + return *Parser_; + } + +private: + bool AdvanceRow() { + while (!Parser_ || !Parser_->TryNextRow()) { + auto next = Producer_->TryGetNextResultSet(); + if (!next) { + Parser_.reset(); + return false; + } + Parser_.emplace(std::move(*next)); + } + return true; + } + + class IResultSetProducer { + public: + virtual ~IResultSetProducer() = default; + virtual std::optional TryGetNextResultSet() = 0; + }; + + class TOneShotProducer : public IResultSetProducer { + public: + explicit TOneShotProducer(TResultSet&& resultSet) + : ResultSet_(std::move(resultSet)) + {} + + std::optional TryGetNextResultSet() override { + if (!ResultSet_) { + return std::nullopt; + } + std::optional result; + result.swap(ResultSet_); + return result; + } + + private: + std::optional ResultSet_; + }; + + class TEmptyProducer : public IResultSetProducer { + public: + std::optional TryGetNextResultSet() override { + return std::nullopt; + } + }; + + template + class TStreamProducer : public IResultSetProducer { + public: + explicit TStreamProducer(Iterator&& iterator) + : Iterator_(std::move(iterator)) + {} + + std::optional TryGetNextResultSet() override { + return NRowRangesDetail::DrainStreamIterator(Iterator_); + } + + private: + Iterator Iterator_; + }; + + std::unique_ptr Producer_; + std::optional Parser_; + bool Started_ = false; +}; + +TRowParser::TRowParser(TResultSetParser& parser) + : Parser_(parser) +{} + +class TRowParserHolder { +public: + explicit TRowParserHolder(TResultSetParser& parser) + : Parser_(parser) + {} + + TRowParser& Get() { + return Parser_; + } + +private: + TRowParser Parser_; +}; + +size_t TRowParser::ColumnsCount() const { + return Parser_.ColumnsCount(); +} + +size_t TRowParser::RowsCount() const { + return Parser_.RowsCount(); +} + +ssize_t TRowParser::ColumnIndex(const std::string& columnName) { + return Parser_.ColumnIndex(columnName); +} + +TValueParser& TRowParser::ColumnParser(size_t columnIndex) { + return Parser_.ColumnParser(columnIndex); +} + +TValueParser& TRowParser::ColumnParser(const std::string& columnName) { + return Parser_.ColumnParser(columnName); +} + +TValue TRowParser::GetValue(size_t columnIndex) const { + return Parser_.GetValue(columnIndex); +} + +TValue TRowParser::GetValue(const std::string& columnName) const { + return Parser_.GetValue(columnName); +} + +class TRowIterator::TImpl { +public: + explicit TImpl(std::shared_ptr range) + : Range_(std::move(range)) + { + ResetRowView(); + } + + bool IsAtEnd() const noexcept { + return !Range_ || !Range_->HasCurrent(); + } + + void Advance() { + RowView_.reset(); + if (Range_) { + Range_->Advance(); + } + ResetRowView(); + } + + TRowParser& Current() { + Y_ABORT_UNLESS(RowView_.has_value(), "TRowIterator dereference at end"); + return RowView_->Get(); + } + +private: + void ResetRowView() { + if (Range_ && Range_->HasCurrent()) { + RowView_.emplace(Range_->CurrentParser()); + } + } + + std::shared_ptr Range_; + std::optional RowView_; +}; + +TRowRange::TRowRange(TResultSet&& resultSet) + : Impl_(std::make_shared(std::move(resultSet))) +{} + +TRowRange::TRowRange(NTable::TDataQueryResult&& result) { + ThrowIfUnsuccessful(result); + auto sets = std::move(result).ExtractResultSets(); + if (sets.size() > 1) { + ythrow yexception() << "multiple queries in one range is not allowed"; + } + if (sets.empty()) { + Impl_ = std::make_shared(TRowRange::TImpl::TEmptyTag{}); + } else { + Impl_ = std::make_shared(std::move(sets[0])); + } +} + +TRowRange::TRowRange(NQuery::TExecuteQueryIterator&& iterator) { + ThrowIfUnsuccessful(iterator); + Impl_ = std::make_shared(std::move(iterator)); +} + +TRowRange::TRowRange(NTable::TScanQueryPartIterator&& iterator) { + ThrowIfUnsuccessful(iterator); + Impl_ = std::make_shared(std::move(iterator)); +} + +TRowRange::TRowRange(NTable::TTablePartIterator&& iterator) { + ThrowIfUnsuccessful(iterator); + Impl_ = std::make_shared(std::move(iterator)); +} + +TRowRange::TRowRange(TRowRange&&) noexcept = default; +TRowRange& TRowRange::operator=(TRowRange&&) noexcept = default; +TRowRange::~TRowRange() = default; + +TRowIterator TRowRange::begin() { + return BeginIterator(Impl_); +} + +TRowIterEnd TRowRange::end() const noexcept { + return TRowIterEnd{}; +} + +TRowIterator TRowRange::BeginIterator(const std::shared_ptr& state) { + if (state) { + state->Start(); + } + return TRowIterator(std::make_unique(state)); +} + +TRowIterator::TRowIterator(std::unique_ptr impl) + : Impl_(std::move(impl)) +{} + +TRowIterator::TRowIterator(TRowIterator&&) noexcept = default; +TRowIterator& TRowIterator::operator=(TRowIterator&&) noexcept = default; +TRowIterator::~TRowIterator() = default; + +bool TRowIterator::operator==(const TRowIterEnd&) const noexcept { + return Impl_->IsAtEnd(); +} + +bool TRowIterator::operator!=(const TRowIterEnd& end) const noexcept { + return !(*this == end); +} + +TRowParser& TRowIterator::operator*() const { + return Impl_->Current(); +} + +TRowParser* TRowIterator::operator->() const { + return &Impl_->Current(); +} + +TRowIterator& TRowIterator::operator++() { + Impl_->Advance(); + return *this; +} + +} // namespace NYdb::inline V3 diff --git a/src/client/row_ranges/rows_stream_drain.h b/src/client/row_ranges/rows_stream_drain.h new file mode 100644 index 00000000000..1bea662cbeb --- /dev/null +++ b/src/client/row_ranges/rows_stream_drain.h @@ -0,0 +1,57 @@ +#pragma once + +#include +#include +#include +#include + +#include + +#include + +#include + +namespace NYdb::inline V3::NRowRangesDetail { + +inline std::optional TryExtractStreamPart(NQuery::TExecuteQueryPart& part) { + if (!part.HasResultSet()) { + return std::nullopt; + } + if (part.GetResultSetIndex() > 0) { + ythrow yexception() << "multiple queries in one range is not allowed"; + } + return std::make_optional(part.ExtractResultSet()); +} + +inline std::optional TryExtractStreamPart(NTable::TScanQueryPart& part) { + if (!part.HasResultSet()) { + return std::nullopt; + } + return std::make_optional(part.ExtractResultSet()); +} + +inline std::optional TryExtractStreamPart(NTable::TReadTableResultPart& part) { + return std::make_optional(part.ExtractPart()); +} + +//! Drains the next result-set chunk from a stream iterator (sync over async). +//! Returns std::nullopt on end-of-stream (EOS). +//! For ExecuteQuery parts, rejects result_set_index > 0 via TryExtractStreamPart. +template +std::optional DrainStreamIterator(Iterator& iterator) { + for (;;) { + auto part = iterator.ReadNext().ExtractValueSync(); + if (!part.IsSuccess()) { + if (part.EOS()) { + return std::nullopt; + } + TStatus status(std::move(part)); + throw NStatusHelpers::TYdbRangeErrorException(status) << status; + } + if (auto extracted = TryExtractStreamPart(part)) { + return extracted; + } + } +} + +} // namespace NYdb::inline V3::NRowRangesDetail diff --git a/src/client/table/CMakeLists.txt b/src/client/table/CMakeLists.txt index 18e607cf6a0..794ac4b1e00 100644 --- a/src/client/table/CMakeLists.txt +++ b/src/client/table/CMakeLists.txt @@ -14,6 +14,7 @@ target_link_libraries(client-ydb_table PUBLIC client-ydb_params client-ydb_proto client-ydb_result + client-ydb_row_ranges client-ydb_scheme client-ydb_table-impl client-ydb_table-query_stats diff --git a/src/client/table/impl/table_client.cpp b/src/client/table/impl/table_client.cpp index 8cd6b6a970d..9a6970699f8 100644 --- a/src/client/table/impl/table_client.cpp +++ b/src/client/table/impl/table_client.cpp @@ -1,5 +1,8 @@ #include "table_client.h" +#include +#include + namespace NYdb::inline V3 { namespace NTable { @@ -1155,6 +1158,10 @@ void TTableClient::TImpl::DeleteSession(TKqpSessionCommon* sessionImpl) { delete sessionImpl; } +void TTableClient::TImpl::PessimizeNode(std::uint64_t nodeId) { + DbDriverState_->EndpointPool.BanNodeId(nodeId); +} + ui32 TTableClient::TImpl::GetSessionRetryLimit() const { return Settings_.SessionPoolSettings_.RetryLimit_; } @@ -1171,8 +1178,13 @@ void TTableClient::TImpl::SetStatCollector(const NSdkStats::TStatCollector::TCli TAsyncBulkUpsertResult TTableClient::TImpl::BulkUpsert(const std::string& table, TValue&& rows, const TBulkUpsertSettings& settings) { Ydb::Table::BulkUpsertRequest* request = nullptr; std::unique_ptr holder; + std::shared_ptr retryHolder; - if (settings.Arena_) { + if (settings.RetryRowsState_) { + retryHolder = std::make_shared( + MakeOperationRequest(settings)); + request = retryHolder.get(); + } else if (settings.Arena_) { request = MakeOperationRequestOnArena(settings, settings.Arena_); } else { holder = std::make_unique(MakeOperationRequest(settings)); @@ -1197,14 +1209,23 @@ TAsyncBulkUpsertResult TTableClient::TImpl::BulkUpsert(const std::string& table, auto obs = MakeObservation("BulkUpsert"); auto promise = NewPromise(); - auto extractor = [promise, obs](google::protobuf::Any* any, TPlainStatus status) mutable { + auto retryState = settings.RetryRowsState_; + + auto extractor = [promise, obs, retryState, retryHolder]( + google::protobuf::Any* any, TPlainStatus status) mutable + { Y_UNUSED(any); + if (retryState && retryHolder && !retryState->HasBackup() + && !status.Ok() && NRetry::ShouldRetryStatus(status.Status, retryState->Settings())) + { + retryState->CreateBackup(*retryHolder); + } obs->End(status.Status, status.Endpoint); TBulkUpsertResult val(TStatus(std::move(status))); promise.SetValue(std::move(val)); }; - if (settings.Arena_) { + if (settings.Arena_ || retryHolder) { Connections_->RunDeferred( request, extractor, @@ -1404,6 +1425,9 @@ void TTableClient::TImpl::SetTxSettings(const TTxSettings& txSettings, Ydb::Tabl case TTxSettings::TS_SNAPSHOT_RW: proto->mutable_snapshot_read_write(); break; + case TTxSettings::TS_STRICT_SERIALIZABLE_RW: + proto->mutable_strict_serializable_read_write(); + break; default: throw TContractViolation("Unexpected transaction mode."); } diff --git a/src/client/table/impl/table_client.h b/src/client/table/impl/table_client.h index 8bddced9b79..c9d77b9afc6 100644 --- a/src/client/table/impl/table_client.h +++ b/src/client/table/impl/table_client.h @@ -140,6 +140,7 @@ class TTableClient::TImpl: public TClientImplCommon, public bool ReturnSession(TKqpSessionCommon* sessionImpl) override; void DeleteSession(TKqpSessionCommon* sessionImpl) override; + void PessimizeNode(std::uint64_t nodeId) override; ui32 GetSessionRetryLimit() const; void SetStatCollector(const NSdkStats::TStatCollector::TClientStatCollector& collector); @@ -163,7 +164,6 @@ class TTableClient::TImpl: public TClientImplCommon, public , const std::shared_ptr& parent = nullptr ); -public: TClientSettings Settings_; private: diff --git a/src/client/table/table.cpp b/src/client/table/table.cpp index 63619124137..ff46e22adda 100644 --- a/src/client/table/table.cpp +++ b/src/client/table/table.cpp @@ -6,11 +6,14 @@ #include #include #include +#include #include +#include #undef INCLUDE_YDB_INTERNAL_H #include #include +#include #include #include #include @@ -223,6 +226,24 @@ const TCompactionOperation::TMetadata& TCompactionOperation::Metadata() const { return Metadata_; } +TAnalyzeOperation::TAnalyzeOperation(TStatus &&status, Ydb::Operations::Operation &&operation) + : TOperation(std::move(status), std::move(operation)) +{ + Ydb::Table::AnalyzeMetadata metadata; + GetProto().metadata().UnpackTo(&metadata); + Metadata_.State = static_cast(metadata.state()); + Metadata_.Progress = metadata.progress(); + Metadata_.Paths.assign(metadata.paths().begin(), metadata.paths().end()); + Metadata_.InProgressPaths.assign(metadata.in_progress_paths().begin(), + metadata.in_progress_paths().end()); + Metadata_.DonePaths.assign(metadata.done_paths().begin(), + metadata.done_paths().end()); +} + +const TAnalyzeOperation::TMetadata& TAnalyzeOperation::Metadata() const { + return Metadata_; +} + //////////////////////////////////////////////////////////////////////////////// class TPartitioningSettings::TImpl { @@ -1575,6 +1596,18 @@ int64_t TTableClient::GetCurrentPoolSize() const { return Impl_->GetCurrentPoolSize(); } +namespace { +thread_local bool TableClientInRetryOperationContext = false; +} // namespace + +bool TTableClient::GetInRetryOperationContext() const { + return TableClientInRetryOperationContext; +} + +void TTableClient::SetInRetryOperationContext(bool value) { + TableClientInRetryOperationContext = value; +} + TTableBuilder TTableClient::GetTableBuilder() { return TTableBuilder(); } @@ -1614,19 +1647,95 @@ NThreading::TFuture TTableClient::Stop() { TAsyncBulkUpsertResult TTableClient::BulkUpsert(const std::string& table, TValue&& rows, const TBulkUpsertSettings& settings) { - return Impl_->BulkUpsert(table, std::move(rows), settings); + const auto retrySettings = NRetry::ResolveRetrySettings( + Impl_->Settings_.RetrySettings_, + settings.RetrySettings_, + settings.ClientTimeout_, + NRetry::ERetryIdempotentDefault::True); + if (!NRetry::IsRetryEnabled(retrySettings) || GetInRetryOperationContext()) { + return Impl_->BulkUpsert(table, std::move(rows), settings); + } + + auto state = std::make_shared(retrySettings); + auto opSettings = settings; + opSettings.RetryRowsState_ = state; + const auto startedAt = TInstant::Now(); + + return Impl_->BulkUpsert(table, std::move(rows), opSettings).Apply( + [this, table, settings, retrySettings, state, startedAt](const TAsyncBulkUpsertResult& f) { + const auto result = f.GetValue(); + if (result.IsSuccess() + || !NRetry::ShouldRetryStatus(result.GetStatus(), retrySettings)) { + return NThreading::MakeFuture(result); + } + Y_ABORT_UNLESS(state->HasBackup()); + + const auto elapsed = TInstant::Now() - startedAt; + if (retrySettings.MaxTimeout_ != TDuration::Max() && elapsed >= retrySettings.MaxTimeout_) { + return NThreading::MakeFuture(result); + } + + auto remaining = retrySettings; + remaining.MaxRetries(retrySettings.MaxRetries_ - 1); + if (retrySettings.MaxTimeout_ != TDuration::Max()) { + remaining.MaxTimeout(retrySettings.MaxTimeout_ - elapsed); + } + + return NRetry::RunUnaryWithRetry(*this, remaining, + [this, table, state, settings](TDuration timeout) { + auto op = settings; + op.RetryRowsState_.reset(); + if (timeout != TDuration::Max()) { + op.ClientTimeout(timeout); + } + return Impl_->BulkUpsert(table, state->GetBackupCopy(), op); + }); + }); } TAsyncBulkUpsertResult TTableClient::BulkUpsert(const std::string& table, EDataFormat format, const std::string& data, const std::string& schema, const TBulkUpsertSettings& settings) { - return Impl_->BulkUpsert(table, format, data, schema, settings); + const auto retrySettings = NRetry::ResolveRetrySettings( + Impl_->Settings_.RetrySettings_, + settings.RetrySettings_, + settings.ClientTimeout_, + NRetry::ERetryIdempotentDefault::True); + if (!NRetry::IsRetryEnabled(retrySettings) || GetInRetryOperationContext()) { + return Impl_->BulkUpsert(table, format, data, schema, settings); + } + + return NRetry::RunUnaryWithRetry(*this, retrySettings, + [this, table, format, data, schema, settings](TDuration timeout) { + auto opSettings = settings; + if (timeout != TDuration::Max()) { + opSettings.ClientTimeout(timeout); + } + return Impl_->BulkUpsert(table, format, data, schema, opSettings); + }); } TAsyncReadRowsResult TTableClient::ReadRows(const std::string& table, TValue&& rows, const std::vector& columns, const TReadRowsSettings& settings) { - return Impl_->ReadRows(table, std::move(rows), columns, settings); + const auto retrySettings = NRetry::ResolveRetrySettings( + Impl_->Settings_.RetrySettings_, + settings.RetrySettings_, + settings.ClientTimeout_, + NRetry::ERetryIdempotentDefault::True); + if (!NRetry::IsRetryEnabled(retrySettings) || GetInRetryOperationContext()) { + return Impl_->ReadRows(table, std::move(rows), columns, settings); + } + TValue keysCopy = std::move(rows); + + return NRetry::RunUnaryWithRetry(*this, retrySettings, + [this, table, keysCopy = std::move(keysCopy), columns, settings](TDuration timeout) { + auto opSettings = settings; + if (timeout != TDuration::Max()) { + opSettings.ClientTimeout(timeout); + } + return Impl_->ReadRows(table, TValue{keysCopy}, columns, opSettings); + }); } TAsyncScanQueryPartIterator TTableClient::StreamExecuteScanQuery(const std::string& query, const TParams& params, @@ -4123,6 +4232,11 @@ TBulkUpsertResult::TBulkUpsertResult(TStatus&& status) : TStatus(std::move(status)) {} +TReadRowsResult::TReadRowsResult(TStatus&& status) + : TStatus(std::move(status)) + , ResultSet(Ydb::ResultSet{}) +{} + TReadRowsResult::TReadRowsResult(TStatus&& status, TResultSet&& resultSet) : TStatus(std::move(status)) , ResultSet(std::move(resultSet)) diff --git a/src/client/topic/codecs/CMakeLists.txt b/src/client/topic/codecs/CMakeLists.txt index e8069c8debd..aaa31c89161 100644 --- a/src/client/topic/codecs/CMakeLists.txt +++ b/src/client/topic/codecs/CMakeLists.txt @@ -3,6 +3,7 @@ _ydb_sdk_add_library(client-ydb_topic-codecs) target_link_libraries(client-ydb_topic-codecs PUBLIC yutil streams-zstd + library-kafka api-grpc-draft api-grpc api-protos diff --git a/src/client/topic/codecs/codecs.cpp b/src/client/topic/codecs/codecs.cpp index 54713252894..b8a1c21e56c 100644 --- a/src/client/topic/codecs/codecs.cpp +++ b/src/client/topic/codecs/codecs.cpp @@ -1,5 +1,9 @@ #include +#include +#include +#include + #include #include @@ -29,7 +33,9 @@ class TZstdToStringCompressor: private TEmbedPolicy, public TZstd } -std::string TGzipCodec::Decompress(const std::string& data) const { +namespace { + +std::string DecompressGzipData(const std::string& data) { TMemoryInput input(data.data(), data.size()); TString result; TStringOutput resultOutput(result); @@ -38,11 +44,7 @@ std::string TGzipCodec::Decompress(const std::string& data) const { return result; } -std::unique_ptr TGzipCodec::CreateCoder(TBuffer& result, int quality) const { - return std::make_unique(result, ZLib::GZip, quality >= 0 ? quality : 6); -} - -std::string TZstdCodec::Decompress(const std::string& data) const { +std::string DecompressZstdData(const std::string& data) { TMemoryInput input(data.data(), data.size()); TString result; TStringOutput resultOutput(result); @@ -51,6 +53,36 @@ std::string TZstdCodec::Decompress(const std::string& data) const { return result; } +std::string KafkaBytesToString(const NKafka::TKafkaBytes& bytes) { + if (!bytes) { + return {}; + } + return std::string(bytes->data(), bytes->size()); +} + +TDecompressionResult MakeSingleMessageResult(std::string data) { + TDecompressionResult result; + result.Messages.push_back(TDecompressedMessage{ + .Data = std::move(data), + .Meta = std::nullopt, + }); + return result; +} + +} // namespace + +std::string TGzipCodec::Decompress(const std::string& data) const { + return DecompressGzipData(data); +} + +std::unique_ptr TGzipCodec::CreateCoder(TBuffer& result, int quality) const { + return std::make_unique(result, ZLib::GZip, quality >= 0 ? quality : 6); +} + +std::string TZstdCodec::Decompress(const std::string& data) const { + return DecompressZstdData(data); +} + std::unique_ptr TZstdCodec::CreateCoder(TBuffer& result, int quality) const { return std::make_unique(result, quality); } @@ -63,11 +95,125 @@ std::unique_ptr TUnsupportedCodec::CreateCoder(TBuffer&, int) con throw yexception() << "use of unsupported codec"; } +TDecompressionResult ICodec::DecompressData(const std::string& data) const { + return MakeSingleMessageResult(Decompress(data)); +} + +void ICodec::CompressWriteBlock(TWriteBlockCompression& ctx) const { + TBuffer compressedData; + std::unique_ptr coder = CreateCoder(compressedData, ctx.CompressionLevel); + for (auto& buffer : ctx.Payloads) { + coder->Write(buffer.data(), buffer.size()); + } + coder->Finish(); + Y_ABORT_UNLESS(!compressedData.Empty()); + ctx.Data = std::move(compressedData); + ctx.Compressed = true; + ctx.CodecID = static_cast(ctx.Codec); +} + +std::string TKafkaBatchCodec::Decompress(const std::string& data) const { + return TakeFirstDecompressedMessage(DecompressData(data)); +} + +TDecompressionResult TKafkaBatchCodec::DecompressData(const std::string& data) const { + using namespace NKafka; + + TDecompressionResult result; + const TKafkaRecordBatch kafkaBatch = ReadKafkaRecordBatch(data); + result.BatchBaseOffset = kafkaBatch.BaseOffset; + result.BatchBaseSequence = 0; + result.BatchBaseTimestampMs = kafkaBatch.BaseTimestamp; + result.Messages.reserve(kafkaBatch.Records.size()); + + for (size_t i = 0; i < kafkaBatch.Records.size(); ++i) { + const auto& record = kafkaBatch.Records[i]; + TDecompressedMessageMeta meta{ + .OffsetDelta = static_cast(record.OffsetDelta), + .SequenceDelta = static_cast(GetRecordSeqNo(kafkaBatch, i, record)), + .TimestampDelta = record.TimestampDelta, + }; + result.Messages.push_back(TDecompressedMessage{ + .Data = KafkaBytesToString(record.Value), + .Meta = std::move(meta), + }); + } + + return result; +} + +std::unique_ptr TKafkaBatchCodec::CreateCoder(TBuffer&, int) const { + throw yexception() << "use of unsupported codec"; +} + +namespace { + +NKafka::ECompressionType ToKafkaBatchInnerCompression(std::optional batchInnerCodec) { + if (!batchInnerCodec) { + return NKafka::ECompressionType::NONE; + } + switch (*batchInnerCodec) { + case ECodec::GZIP: + return NKafka::ECompressionType::GZIP; + case ECodec::ZSTD: + return NKafka::ECompressionType::ZSTD; + default: + ythrow yexception() << "unsupported batch inner codec: " << static_cast(*batchInnerCodec); + } +} + +} // namespace + +void TKafkaBatchCodec::CompressWriteBlock(TWriteBlockCompression& ctx) const { + if (ctx.Payloads.size() == 1) { + return; + } + + using namespace NKafka; + + Y_ABORT_UNLESS(ctx.Payloads.size() == ctx.CreatedAt.size()); + Y_ABORT_UNLESS(!ctx.Payloads.empty()); + + TKafkaRecordBatch kafkaBatch; + kafkaBatch.Magic = 2; + kafkaBatch.ProducerId = 0; + kafkaBatch.ProducerEpoch = 0; + kafkaBatch.BaseSequence = static_cast(ctx.BaseSequence); + kafkaBatch.Attributes = static_cast( + ToKafkaBatchInnerCompression(ctx.BatchInnerCodec)); + + const i64 baseTimestamp = static_cast(ctx.CreatedAt.front().MilliSeconds()); + kafkaBatch.BaseTimestamp = baseTimestamp; + kafkaBatch.MaxTimestamp = baseTimestamp; + kafkaBatch.LastOffsetDelta = static_cast(ctx.Payloads.size() - 1); + + for (size_t i = 0; i < ctx.Payloads.size(); ++i) { + TKafkaRecord record; + record.OffsetDelta = static_cast(i); + record.TimestampDelta = ctx.CreatedAt[i].MilliSeconds() - baseTimestamp; + kafkaBatch.MaxTimestamp = Max(kafkaBatch.MaxTimestamp, static_cast(ctx.CreatedAt[i].MilliSeconds())); + record.Value = TKafkaRawBytes(ctx.Payloads[i].data(), ctx.Payloads[i].size()); + record.Length = record.Size(2) - NKafka::NPrivate::SizeOfVarint(0); + kafkaBatch.Records.push_back(std::move(record)); + } + + kafkaBatch.BatchLength = kafkaBatch.Size(2) + - sizeof(TKafkaRecordBatch::BaseOffsetMeta::Type) + - sizeof(TKafkaRecordBatch::BatchLengthMeta::Type); + + const TString serialized = WriteKafkaRecordBatch(kafkaBatch); + ctx.Data = TBuffer(serialized.data(), serialized.size()); + ctx.Payloads.assign(1, std::string_view(ctx.Data.data(), ctx.Data.size())); + ctx.CodecID = static_cast(Ydb::Topic::CODEC_KAFKA_BATCH); + ctx.Compressed = true; +} + class TCommonCodecsProvider { public: TCommonCodecsProvider() { TCodecMap::GetTheCodecMap().Set((uint32_t)ECodec::GZIP, std::make_unique()); TCodecMap::GetTheCodecMap().Set((uint32_t)ECodec::ZSTD, std::make_unique()); + TCodecMap::GetTheCodecMap().Set((uint32_t)ECodec::KAFKA_BATCH, std::make_unique()); } }; diff --git a/src/client/topic/impl/common.h b/src/client/topic/impl/common.h index f3430c9020c..89607fc4007 100644 --- a/src/client/topic/impl/common.h +++ b/src/client/topic/impl/common.h @@ -2,13 +2,24 @@ #include #include +#include +#include #include +#include +#include #include +#include + #include #include +#include +#include +#include +#include +#include namespace NYdb::inline V3::NTopic { @@ -21,6 +32,182 @@ NYdb::NIssue::TIssues MakeIssueWithSubIssues(const std::string& description, con std::string IssuesSingleLineString(const NYdb::NIssue::TIssues& issues); +namespace NWriteSessionGrpc { + +inline constexpr std::string_view PARTITION_KEY_META_KEY = "__partition_key"; + +inline size_t GetMaxGrpcMessageSize() { + return 120_MB; +} + +using TWireFormatLite = google::protobuf::internal::WireFormatLite; + +inline size_t ProtoInt32FieldSize(ui32 fieldNumber, i32 value) { + if (value == 0) { + return 0; + } + return TWireFormatLite::TagSize(fieldNumber, TWireFormatLite::TYPE_INT32) + + TWireFormatLite::Int32Size(value); +} + +inline size_t ProtoInt64FieldSize(ui32 fieldNumber, i64 value) { + if (value == 0) { + return 0; + } + return TWireFormatLite::TagSize(fieldNumber, TWireFormatLite::TYPE_INT64) + + TWireFormatLite::Int64Size(value); +} + +inline size_t ProtoInt64FieldSizeUpperBound(ui32 fieldNumber) { + static constexpr size_t MaxInt64VarintPayloadSize = 10; + + return TWireFormatLite::TagSize(fieldNumber, TWireFormatLite::TYPE_INT64) + + MaxInt64VarintPayloadSize; +} + +inline size_t ProtoPackedInt64FieldSize(ui32 fieldNumber, size_t dataSize) { + if (dataSize == 0) { + return 0; + } + return TWireFormatLite::TagSize(fieldNumber, TWireFormatLite::TYPE_INT64) + + TWireFormatLite::LengthDelimitedSize(dataSize); +} + +inline size_t ProtoBytesFieldSize(ui32 fieldNumber, size_t size) { + if (size == 0) { + return 0; + } + return TWireFormatLite::TagSize(fieldNumber, TWireFormatLite::TYPE_BYTES) + + TWireFormatLite::LengthDelimitedSize(size); +} + +inline size_t ProtoStringFieldSize(ui32 fieldNumber, size_t size) { + if (size == 0) { + return 0; + } + return TWireFormatLite::TagSize(fieldNumber, TWireFormatLite::TYPE_STRING) + + TWireFormatLite::LengthDelimitedSize(size); +} + +inline size_t ProtoMessageFieldSize(ui32 fieldNumber, size_t size) { + return TWireFormatLite::TagSize(fieldNumber, TWireFormatLite::TYPE_MESSAGE) + + TWireFormatLite::LengthDelimitedSize(size); +} + +class TRequestSizeLimiter { +public: + explicit TRequestSizeLimiter(ui32 envelopeFieldNumber, size_t maxSize = GetMaxGrpcMessageSize()) + : EnvelopeFieldNumber(envelopeFieldNumber) + , MaxSize(maxSize) + { + } + + bool Empty() const { + return BodySize == 0; + } + + bool CanAdd(size_t deltaSize) const { + return Empty() || ProtoMessageFieldSize(EnvelopeFieldNumber, BodySize + deltaSize) <= MaxSize; + } + + void Add(size_t deltaSize) { + BodySize += deltaSize; + } + +private: + size_t BodySize = 0; + ui32 EnvelopeFieldNumber; + size_t MaxSize; +}; + +inline size_t ProtoTimestampFieldSize(ui32 fieldNumber, TInstant timestamp) { + const ui64 milliseconds = timestamp.MilliSeconds(); + const i64 seconds = milliseconds / 1000; + const i32 nanos = (milliseconds % 1000) * 1000000; + const size_t timestampSize = ProtoInt64FieldSize(1, seconds) + + ProtoInt32FieldSize(2, nanos); + return ProtoMessageFieldSize(fieldNumber, timestampSize); +} + +inline size_t ProtoMetadataItemFieldSize(ui32 fieldNumber, const std::pair& item) { + const size_t itemSize = ProtoStringFieldSize(1, item.first.size()) + + ProtoBytesFieldSize(2, item.second.size()); + return ProtoMessageFieldSize(fieldNumber, itemSize); +} + +inline size_t ProtoTransactionIdentityFieldSize(ui32 fieldNumber, const std::optional& tx) { + if (!tx) { + return 0; + } + const size_t txSize = ProtoStringFieldSize(1, tx->TxId.size()) + + ProtoStringFieldSize(2, tx->SessionId.size()); + return ProtoMessageFieldSize(fieldNumber, txSize); +} + +inline size_t ProtoTopicMessageDataFieldSize( + TInstant createdAt, + size_t dataSize, + size_t uncompressedSize, + size_t metadataSize) { + const size_t messageDataSize = ProtoInt64FieldSizeUpperBound(1) + + ProtoTimestampFieldSize(2, createdAt) + + ProtoBytesFieldSize(3, dataSize) + + ProtoInt64FieldSize(4, uncompressedSize) + + metadataSize; + return ProtoMessageFieldSize(1, messageDataSize); +} + +template +size_t EstimateTopicWriteRequestBlockSize( + const TBlock& block, + const TOriginalMessages& originalMessages, + bool includeRequestFields) { + Y_ABORT_UNLESS(!originalMessages.empty()); + + size_t size = 0; + if (includeRequestFields) { + size += ProtoInt32FieldSize(2, static_cast(block.CodecID)); + size += ProtoTransactionIdentityFieldSize(3, originalMessages.front().Tx); + } + + if (block.MessageCount > 1) { + size_t metadataSize = 0; + auto message = originalMessages.begin(); + const TInstant firstMessageCreatedAt = message->CreatedAt; + for (const auto& item : message->MessageMeta) { + metadataSize += ProtoMetadataItemFieldSize(7, item); + } + ++message; + for (size_t i = 1; i < block.MessageCount; ++i) { + Y_ABORT_UNLESS(message != originalMessages.end()); + for (const auto& item : message->MessageMeta) { + if (item.first == PARTITION_KEY_META_KEY) { + metadataSize += ProtoMetadataItemFieldSize(7, item); + } + } + ++message; + } + return size + ProtoTopicMessageDataFieldSize(firstMessageCreatedAt, block.Data.size(), block.OriginalSize, metadataSize); + } + + const auto& message = originalMessages.front(); + size_t metadataSize = 0; + for (const auto& item : message.MessageMeta) { + metadataSize += ProtoMetadataItemFieldSize(7, item); + } + + if (block.Compressed) { + return size + ProtoTopicMessageDataFieldSize(message.CreatedAt, block.Data.size(), block.OriginalSize, metadataSize); + } + + for (const auto& buffer : block.OriginalDataRefs) { + size += ProtoTopicMessageDataFieldSize(message.CreatedAt, buffer.size(), block.OriginalSize, metadataSize); + } + return size; +} + +} // namespace NWriteSessionGrpc + template size_t CalcDataSize(const typename TEvent::TEvent& event) { constexpr bool UseMigrationProtocol = !std::is_same_v; diff --git a/src/client/topic/impl/offsets_collector.cpp b/src/client/topic/impl/offsets_collector.cpp index b846ccf703c..11c6a77598e 100644 --- a/src/client/topic/impl/offsets_collector.cpp +++ b/src/client/topic/impl/offsets_collector.cpp @@ -59,14 +59,14 @@ void TOffsetsCollector::CollectOffsets(const TReadSessionEvent::TDataReceivedEve if (event.HasCompressedMessages()) { for (auto& message : event.GetCompressedMessages()) { uint64_t offset = message.GetOffset(); - Ranges[topicPath][partitionId].InsertInterval(offset, offset + 1); + Ranges[topicPath][partitionId].InsertInterval(offset, offset + message.GetLogicalMessageCount()); } } else { for (auto& message : event.GetMessages()) { uint64_t offset = message.GetOffset(); - Ranges[topicPath][partitionId].InsertInterval(offset, offset + 1); + Ranges[topicPath][partitionId].InsertInterval(offset, offset + message.GetLogicalMessageCount()); } } } -} \ No newline at end of file +} diff --git a/src/client/topic/impl/read_session_event.cpp b/src/client/topic/impl/read_session_event.cpp index 7c1871728ac..b935e5daa19 100644 --- a/src/client/topic/impl/read_session_event.cpp +++ b/src/client/topic/impl/read_session_event.cpp @@ -26,10 +26,10 @@ using TPartitionSessionClosedEvent = TReadSessionEvent::TPartitionSessionClosedE std::pair GetMessageOffsetRange(const TDataReceivedEvent& dataReceivedEvent, uint64_t index) { if (dataReceivedEvent.HasCompressedMessages()) { const auto& msg = dataReceivedEvent.GetCompressedMessages()[index]; - return {msg.GetOffset(), msg.GetOffset() + 1}; + return {msg.GetOffset(), msg.GetOffset() + msg.GetLogicalMessageCount()}; } const auto& msg = dataReceivedEvent.GetMessages()[index]; - return {msg.GetOffset(), msg.GetOffset() + 1}; + return {msg.GetOffset(), msg.GetOffset() + msg.GetLogicalMessageCount()}; } //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -44,7 +44,8 @@ TMessageInformation::TMessageInformation( TWriteSessionMeta::TPtr meta, TMessageMeta::TPtr messageMeta, uint64_t uncompressedSize, - std::string messageGroupId + std::string messageGroupId, + uint64_t logicalMessageCount ) : Offset(offset) , ProducerId(producerId) @@ -55,6 +56,7 @@ TMessageInformation::TMessageInformation( , MessageMeta(messageMeta) , UncompressedSize(uncompressedSize) , MessageGroupId(std::move(messageGroupId)) + , LogicalMessageCount(logicalMessageCount) {} //////////////////////////////////////////////////////////////////////////////////////////////////////////////////////// @@ -96,6 +98,10 @@ uint64_t TMessageBase::GetOffset() const { return Information.Offset; } +uint64_t TMessageBase::GetLogicalMessageCount() const { + return Information.LogicalMessageCount; +} + const std::string& TMessageBase::GetProducerId() const { return Information.ProducerId; } @@ -139,6 +145,7 @@ void TPrintable::DebugString(TStringBuilder& ret, bool printData) } ret << " Information: {" << " Offset: " << self->GetOffset() + << " LogicalMessageCount: " << self->GetLogicalMessageCount() << " ProducerId: \"" << self->GetProducerId() << "\"" << " SeqNo: " << self->GetSeqNo() << " CreateTime: " << self->GetCreateTime() @@ -192,7 +199,7 @@ bool TMessage::HasException() const { void TMessage::Commit() { static_cast(PartitionSession.Get()) - ->Commit(Information.Offset, Information.Offset + 1); + ->Commit(Information.Offset, Information.Offset + Information.LogicalMessageCount); } template<> @@ -226,7 +233,7 @@ uint64_t TCompressedMessage::GetUncompressedSize() const { void TCompressedMessage::Commit() { static_cast(PartitionSession.Get()) - ->Commit(Information.Offset, Information.Offset + 1); + ->Commit(Information.Offset, Information.Offset + Information.LogicalMessageCount); } template<> @@ -315,10 +322,10 @@ TStartPartitionSessionEvent::TStartPartitionSessionEvent(TPartitionSession::TPtr , EndOffset(endOffset) { } -void TStartPartitionSessionEvent::Confirm(std::optional readOffset, std::optional commitOffset) { +void TStartPartitionSessionEvent::Confirm(std::optional readOffset, std::optional commitOffset, std::optional maxOffset) { if (PartitionSession) { static_cast(PartitionSession.Get()) - ->ConfirmCreate(readOffset, commitOffset); + ->ConfirmCreate(readOffset, commitOffset, maxOffset); } } diff --git a/src/client/topic/impl/read_session_impl.h b/src/client/topic/impl/read_session_impl.h index 3252c165734..483450d0e3c 100644 --- a/src/client/topic/impl/read_session_impl.h +++ b/src/client/topic/impl/read_session_impl.h @@ -72,7 +72,7 @@ using TASessionClosedEvent = std::conditional_t readOffset, std::optional commitOffset) = 0; + virtual void ConfirmCreate(std::optional readOffset, std::optional commitOffset, std::optional maxOffset) = 0; virtual void ConfirmDestroy() = 0; virtual void ConfirmEnd(std::span childIds) = 0; }; @@ -127,7 +127,7 @@ using TCallbackContextPtr = std::shared_ptr class TUserRetrievedEventsInfoAccumulator { public: - void Add(TDataDecompressionInfoPtr info, i64 decompressedSize); + void Add(TDataDecompressionInfoPtr info, i64 decompressedSize, size_t messagesCount = 1); void OnUserRetrievedEvent() const; private: @@ -231,6 +231,15 @@ template class TDataDecompressionInfo : public std::enable_shared_from_this> { public: using TPtr = std::shared_ptr>; + using TMessage = typename TADataReceivedEvent::TMessage; + using TCompressedMessage = typename TADataReceivedEvent::TCompressedMessage; + + struct TDecompressedData { + std::vector Messages; + std::vector CompressedMessages; + size_t DataSize = 0; + size_t MessagesTaken = 0; + }; TDataDecompressionInfo(const TDataDecompressionInfo&) = default; TDataDecompressionInfo(TDataDecompressionInfo&&) = default; @@ -238,7 +247,8 @@ class TDataDecompressionInfo : public std::enable_shared_from_this&& msg, TCallbackContextPtr cbContext, bool doDecompress, - i64 serverBytesSize = 0 // to increment read request bytes size + i64 serverBytesSize = 0, // to increment read request bytes size + ui64 committedOffset = 0 ); ~TDataDecompressionInfo(); @@ -317,6 +327,14 @@ class TDataDecompressionInfo : public std::enable_shared_from_this> partitionStream, + size_t batch, + size_t message, + size_t& maxByteSize); + + size_t GetPreparedDataSize(size_t batch, size_t message) const; + size_t GetPreparedMessageCount(size_t batch, size_t message) const; + void OnDataDecompressed(i64 sourceSize, i64 estimatedDecompressedSize, i64 decompressedSize, size_t messagesCount); void OnUserRetrievedEvent(i64 decompressedDataSize, size_t messagesCount); void OnTaskCanceled(i64 sourceSize, size_t messagesCount); @@ -366,6 +384,13 @@ class TDataDecompressionInfo : public std::enable_shared_from_this> partitionStream, + size_t batch, + size_t message, + const TDecompressionResult& codecResult); + void PutDecompressedData(size_t batch, + size_t message, + TDecompressedData&& data); private: TPartitionData ServerMessage; @@ -373,8 +398,10 @@ class TDataDecompressionInfo : public std::enable_shared_from_this::TPtr>; TMetadataPtrVector BatchesMeta; std::vector MessagesMeta; + std::vector> DecompressedData; TCallbackContextPtr CbContext; bool DoDecompress; + ui64 CommittedOffset = 0; std::atomic ServerBytesSize = 0; std::atomic SourceDataNotProcessed = 0; std::pair CurrentDecompressingMessage = {0, 0}; // (Batch, Message) @@ -418,13 +445,12 @@ class TDataDecompressionEvent { return !Ready; } - void TakeData(TIntrusivePtr> partitionStream, - std::vector::TMessage>& messages, - std::vector::TCompressedMessage>& compressedMessages, - size_t& maxByteSize, - size_t& dataSize) const; + typename TDataDecompressionInfo::TDecompressedData + TakeData(TIntrusivePtr> partitionStream, + size_t& maxByteSize) const; size_t GetDataSize() const; + size_t GetMessageCount() const; TDataDecompressionInfoPtr GetParent() const { return Parent; @@ -729,7 +755,7 @@ class TPartitionStreamImpl : public TAPartitionStream { void Commit(uint64_t startOffset, uint64_t endOffset) override; void RequestStatus() override; - void ConfirmCreate(std::optional readOffset, std::optional commitOffset) override; + void ConfirmCreate(std::optional readOffset, std::optional commitOffset, std::optional maxOffset) override; void ConfirmDestroy() override; void ConfirmEnd(std::span childIds) override; @@ -901,7 +927,7 @@ class TReadSessionEventsQueue: public TBaseSessionEventsQueue& settings); // Assumes we are under lock. - TReadSessionEventInfo + std::optional> GetEventImpl(size_t& maxByteSize, TUserRetrievedEventsInfoAccumulator& accumulator); @@ -1120,7 +1146,7 @@ class TReadSessionEventsQueue: public TBaseSessionEventsQueue CbContext; }; - TADataReceivedEvent + std::optional> GetDataEventImpl(TIntrusivePtr> stream, size_t& maxByteSize, TUserRetrievedEventsInfoAccumulator& accumulator); // Assumes that we're under lock. @@ -1218,7 +1244,7 @@ class TSingleClusterReadSessionImpl : public TEnableSelfContext* partitionStream, std::optional readOffset, std::optional commitOffset); + void ConfirmPartitionStreamCreate(const TPartitionStreamImpl* partitionStream, std::optional readOffset, std::optional commitOffset, std::optional maxOffset); void ConfirmPartitionStreamDestroy(TPartitionStreamImpl* partitionStream); void ConfirmPartitionStreamEnd(TPartitionStreamImpl* partitionStream, std::span childIds); void RequestPartitionStreamStatus(const TPartitionStreamImpl* partitionStream); diff --git a/src/client/topic/impl/read_session_impl.ipp b/src/client/topic/impl/read_session_impl.ipp index 8a3f7b0eb45..c2a6c38e3be 100644 --- a/src/client/topic/impl/read_session_impl.ipp +++ b/src/client/topic/impl/read_session_impl.ipp @@ -16,6 +16,8 @@ #include +#include + #include #include @@ -31,6 +33,26 @@ namespace NYdb::inline V3::NTopic { +namespace { + +using NKafka::ReadKafkaBatchHeader; + +size_t GetReadMessageCount(const Ydb::Topic::StreamReadMessage_ReadResponse_MessageData& messageData, int32_t codec) { + const auto& dataBytes = messageData.data(); + + switch (codec) { + case Ydb::Topic::CODEC_KAFKA_BATCH: + if (const auto header = NKafka::ReadKafkaBatchHeader(dataBytes)) { + return static_cast(header->RecordsCount); + } + return 1; + default: + return 1; + } +} + +} // namespace + static const bool RangesMode = !std::string{std::getenv("PQ_OFFSET_RANGES_MODE") ? std::getenv("PQ_OFFSET_RANGES_MODE") : ""}.empty(); static const bool ExperimentalDirectRead = !std::string{std::getenv("PQ_EXPERIMENTAL_DIRECT_READ") ? std::getenv("PQ_EXPERIMENTAL_DIRECT_READ") : ""}.empty(); static const bool DecompressEverything = !std::string{std::getenv("PQ_DECOMPRESS_EVERYTHING") ? std::getenv("PQ_DECOMPRESS_EVERYTHING") : ""}.empty(); @@ -78,12 +100,12 @@ void TPartitionStreamImpl::RequestStatus() { } template -void TPartitionStreamImpl::ConfirmCreate(std::optional readOffset, std::optional commitOffset) { +void TPartitionStreamImpl::ConfirmCreate(std::optional readOffset, std::optional commitOffset, std::optional maxOffset) { if (auto sessionShared = CbContext->LockShared()) { if (commitOffset.has_value()) { SetFirstNotReadOffset(commitOffset.value()); } - sessionShared->ConfirmPartitionStreamCreate(this, readOffset, commitOffset); + sessionShared->ConfirmPartitionStreamCreate(this, readOffset, commitOffset, maxOffset); } } @@ -224,7 +246,10 @@ void TRawPartitionStreamEventQueue::DeleteNotReadyTail(TDe ready.push_back(std::move(event)); } else if (isDataEvent) { // We should release memory for this ready event here - accumulator.Add(event.GetDataEvent().GetParent(), event.GetDataEvent().GetDataSize()); + accumulator.Add( + event.GetDataEvent().GetParent(), + event.GetDataEvent().GetDataSize(), + event.GetDataEvent().GetMessageCount()); } } else { hasNonReadyEvents = true; @@ -557,6 +582,8 @@ inline void TSingleClusterReadSessionImpl::InitImpl(TDeferredActionsset_path(TStringType{topic.Path_}); @@ -647,17 +674,21 @@ bool TSingleClusterReadSessionImpl::IsActualPartitionStrea } template -void TSingleClusterReadSessionImpl::ConfirmPartitionStreamCreate(const TPartitionStreamImpl* partitionStream, std::optional readOffset, std::optional commitOffset) { - TStringBuilder commitOffsetLogStr; +void TSingleClusterReadSessionImpl::ConfirmPartitionStreamCreate(const TPartitionStreamImpl* partitionStream, + std::optional readOffset, std::optional commitOffset, std::optional maxOffset) { + TStringBuilder offsetLogStr; if (commitOffset) { - commitOffsetLogStr << ". Commit offset: " << *commitOffset; + offsetLogStr << ". Commit offset: " << *commitOffset; + } + if (maxOffset) { + offsetLogStr << ". Max offset: " << *maxOffset; } LOG_LAZY(Log, TLOG_INFO, GetLogPrefix() << "Confirm partition stream create. Partition stream id: " << GetPartitionStreamId(partitionStream) << ". Cluster: \"" << GetCluster(partitionStream) << "\". Topic: \"" << partitionStream->GetTopicPath() << "\". Partition: " << partitionStream->GetPartitionId() - << ". Read offset: " << readOffset << commitOffsetLogStr + << ". Read offset: " << readOffset << offsetLogStr ); std::lock_guard guard(Lock); @@ -695,6 +726,9 @@ void TSingleClusterReadSessionImpl::ConfirmPartitionStream if (commitOffset) { startRead.set_commit_offset(*commitOffset); } + if (maxOffset) { + startRead.set_max_offset(*maxOffset); + } WriteToProcessorImpl(std::move(req)); @@ -1069,6 +1103,7 @@ inline void TSingleClusterReadSessionImpl::OnReadDoneImpl( return; } const TIntrusivePtr>& partitionStream = partitionStreamIt->second; + Y_ABORT_UNLESS(partitionStream); typename TPartitionCookieMapping::TCookie::TPtr cookie = MakeIntrusive(partitionData.cookie().partition_cookie(), partitionStream); @@ -1149,9 +1184,10 @@ inline void TSingleClusterReadSessionImpl::OnReadDoneImpl( NextPartitionStreamId += PartitionStreamIdStep; // Renew partition stream. - TIntrusivePtr>& currentPartitionStream = - PartitionStreams[partitionStream->GetAssignId()]; - if (currentPartitionStream) { + auto [partitionStreamIt, inserted] = PartitionStreams.emplace(partitionStream->GetAssignId(), partitionStream); + TIntrusivePtr>& currentPartitionStream = partitionStreamIt->second; + if (!inserted) { + Y_ABORT_UNLESS(currentPartitionStream); CookieMapping.RemoveMapping(currentPartitionStream->GetPartitionStreamId()); bool pushRes = EventsQueue->PushEvent( @@ -1163,8 +1199,8 @@ inline void TSingleClusterReadSessionImpl::OnReadDoneImpl( AbortImpl(); return; } + currentPartitionStream = partitionStream; } - currentPartitionStream = partitionStream; // Send event to user. bool pushRes = EventsQueue->PushEvent( @@ -1321,6 +1357,7 @@ inline void TSingleClusterReadSessionImpl::OnReadDoneImpl( } } const TIntrusivePtr>& partitionStream = partitionStreamIt->second; + Y_ABORT_UNLESS(partitionStream); i64 firstOffset = std::numeric_limits::max(); i64 currentOffset = std::numeric_limits::max(); @@ -1329,7 +1366,9 @@ inline void TSingleClusterReadSessionImpl::OnReadDoneImpl( // Validate messages. for (const auto& messageData : batch.message_data()) { // Check offsets continuity. - if (messageData.offset() != desiredOffset) { + // A message may start below the desired offset when a multi-message batch + // is redelivered as a whole blob after restarting from an offset inside it. + if (messageData.offset() > desiredOffset) { bool res = partitionStream->AddToCommitRanges(desiredOffset, messageData.offset(), GetRangesMode()); Y_ABORT_UNLESS(res); } @@ -1338,13 +1377,14 @@ inline void TSingleClusterReadSessionImpl::OnReadDoneImpl( firstOffset = messageData.offset(); } currentOffset = messageData.offset(); - desiredOffset = currentOffset + 1; - partitionStream->UpdateMaxReadOffset(currentOffset); + const ui64 logicalMessageCount = GetReadMessageCount(messageData, batch.codec()); + desiredOffset = Max(desiredOffset, currentOffset + static_cast(logicalMessageCount)); + partitionStream->UpdateMaxReadOffset(currentOffset + static_cast(logicalMessageCount) - 1); const i64 messageSize = static_cast(messageData.data().size()); CompressedDataSize += messageSize; *Settings.Counters_->BytesInflightTotal += messageSize; *Settings.Counters_->BytesInflightCompressed += messageSize; - ++*Settings.Counters_->MessagesInflight; + *Settings.Counters_->MessagesInflight += logicalMessageCount; } } if (firstOffset == std::numeric_limits::max()) { @@ -1357,10 +1397,12 @@ inline void TSingleClusterReadSessionImpl::OnReadDoneImpl( } partitionStream->SetFirstNotReadOffset(desiredOffset); + const ui64 committedOffset = partitionStream->GetMaxCommittedOffset(); auto decompressionInfo = std::make_shared>(std::move(partitionData), SelfContext, Settings.Decompress_, - serverBytesSize); + serverBytesSize, + committedOffset); // TODO (ildar-khisam@): share serverBytesSize between partitions data according to their actual sizes; // for now whole serverBytesSize goes with first (and only) partition data. serverBytesSize = 0; @@ -1380,6 +1422,7 @@ template <> inline void TSingleClusterReadSessionImpl::StopPartitionSessionImpl( TIntrusivePtr> partitionStream, bool graceful, TDeferredActions& deferred ) { + Y_ABORT_UNLESS(Lock.IsLocked()); auto partitionSessionId = partitionStream->GetAssignId(); if (IsDirectRead()) { @@ -1535,12 +1578,14 @@ inline void TSingleClusterReadSessionImpl::OnReadDoneImpl( auto partitionSessionId = msg.partition_session().partition_session_id(); // Renew partition stream. - TIntrusivePtr>& partitionStream = PartitionStreams[partitionSessionId]; - if (partitionStream) { + auto partitionStreamIt = PartitionStreams.find(partitionSessionId); + if (partitionStreamIt != PartitionStreams.end()) { + const TIntrusivePtr>& currentPartitionStream = partitionStreamIt->second; + Y_ABORT_UNLESS(currentPartitionStream); bool pushRes = EventsQueue->PushEvent( - partitionStream, + currentPartitionStream, TReadSessionEvent::TPartitionSessionClosedEvent( - partitionStream, TReadSessionEvent::TPartitionSessionClosedEvent::EReason::Lost), + currentPartitionStream, TReadSessionEvent::TPartitionSessionClosedEvent::EReason::Lost), deferred); if (!pushRes) { @@ -1549,7 +1594,7 @@ inline void TSingleClusterReadSessionImpl::OnReadDoneImpl( } } - partitionStream = MakeIntrusive>( + auto partitionStream = MakeIntrusive>( NextPartitionStreamId, msg.partition_session().path(), ReadSessionId, @@ -1560,6 +1605,7 @@ inline void TSingleClusterReadSessionImpl::OnReadDoneImpl( SelfContext); NextPartitionStreamId += PartitionStreamIdStep; + PartitionStreams.insert_or_assign(partitionSessionId, partitionStream); // Send event to user. bool pushRes = EventsQueue->PushEvent( @@ -2564,17 +2610,22 @@ void TRawPartitionStreamEventQueue::GetDataEventImpl(TIntr auto& event = queue.front().GetDataEvent(); TDataDecompressionInfoPtr parent = event.GetParent(); - size_t size = 0; - event.TakeData(partitionStream, messages, compressedMessages, maxByteSize, size); + auto decompressedData = event.TakeData(partitionStream, maxByteSize); + messages.insert(messages.end(), + std::make_move_iterator(decompressedData.Messages.begin()), + std::make_move_iterator(decompressedData.Messages.end())); + compressedMessages.insert(compressedMessages.end(), + std::make_move_iterator(decompressedData.CompressedMessages.begin()), + std::make_move_iterator(decompressedData.CompressedMessages.end())); queue.pop_front(); - accumulator.Add(parent, size); + accumulator.Add(parent, decompressedData.DataSize, decompressedData.MessagesTaken); } } template -TADataReceivedEvent +std::optional> TReadSessionEventsQueue::GetDataEventImpl(TIntrusivePtr> stream, size_t& maxByteSize, TUserRetrievedEventsInfoAccumulator& accumulator) // Assumes that we're under lock. @@ -2600,19 +2651,21 @@ TReadSessionEventsQueue::GetDataEventImpl(TIntrusivePtr{std::move(messages), std::move(compressedMessages), stream}; } template -TReadSessionEventInfo +std::optional> TReadSessionEventsQueue::GetEventImpl(size_t& maxByteSize, TUserRetrievedEventsInfoAccumulator& accumulator) // Assumes that we're under lock. { Y_ASSERT(TParent::HasEventsImpl()); - if (!TParent::Events.empty()) { + while (!TParent::Events.empty()) { TReadSessionEventInfo& front = TParent::Events.front(); auto partitionStream = front.PartitionStream; @@ -2623,7 +2676,15 @@ TReadSessionEventsQueue::GetEventImpl(size_t& maxByteSize, std::optional::TEvent> event; auto frontCbContext = front.CbContext; if (partitionStream->TopEvent().IsDataEvent()) { - event = GetDataEventImpl(partitionStream, maxByteSize, accumulator); + auto dataEvent = GetDataEventImpl(partitionStream, maxByteSize, accumulator); + if (!dataEvent) { + TParent::RenewWaiterImpl(); + if (maxByteSize == 0) { + return std::nullopt; + } + continue; + } + event = std::move(*dataEvent); } else { event = std::move(partitionStream->TopEvent().GetEvent()); partitionStream->PopEvent(); @@ -2642,12 +2703,14 @@ TReadSessionEventsQueue::GetEventImpl(size_t& maxByteSize, TParent::RenewWaiterImpl(); - return {partitionStream, std::move(frontCbContext), std::move(*event)}; + return TReadSessionEventInfo{partitionStream, std::move(frontCbContext), std::move(*event)}; } - Y_ASSERT(TParent::CloseEvent); + if (!TParent::CloseEvent) { + return std::nullopt; + } - return {*TParent::CloseEvent}; + return TReadSessionEventInfo{*TParent::CloseEvent}; } template @@ -2670,8 +2733,11 @@ TReadSessionEventsQueue::GetEvents(bool block, std::option } while (TParent::HasEventsImpl() && eventInfos.size() < maxCount && maxByteSize > 0) { - TReadSessionEventInfo event = GetEventImpl(maxByteSize, accumulator); - eventInfos.emplace_back(std::move(event)); + auto event = GetEventImpl(maxByteSize, accumulator); + if (!event) { + continue; + } + eventInfos.emplace_back(std::move(*event)); if (eventInfos.back().IsSessionClosedEvent()) { break; } @@ -2827,20 +2893,28 @@ TDataDecompressionInfo::TDataDecompressionInfo( TPartitionData&& msg, TCallbackContextPtr cbContext, bool doDecompress, - i64 serverBytesSize + i64 serverBytesSize, + ui64 committedOffset ) : ServerMessage(std::move(msg)) , CbContext(std::move(cbContext)) , DoDecompress(doDecompress) + , CommittedOffset(committedOffset) , ServerBytesSize(serverBytesSize) { i64 compressedSize = 0; i64 messagesCount = 0; for (const auto& batch : ServerMessage.batches()) { + DecompressedData.emplace_back(); + DecompressedData.back().resize(static_cast(batch.message_data_size())); for (const auto& messageData : batch.message_data()) { compressedSize += messageData.data().size(); - ++messagesCount; + if constexpr (UseMigrationProtocol) { + ++messagesCount; + } else { + messagesCount += static_cast(GetReadMessageCount(messageData, batch.codec())); + } } } @@ -2971,6 +3045,7 @@ template void TDataDecompressionInfo::PlanDecompressionTasks(double averageCompressionRatio, TIntrusivePtr> partitionStream) { constexpr size_t TASK_LIMIT = 512_KB; + Y_ABORT_UNLESS(partitionStream); auto session = CbContext->LockShared(); Y_ASSERT(session); @@ -3035,17 +3110,18 @@ void TDataDecompressionInfo::OnDestroyReadSession() } template -void TDataDecompressionEvent::TakeData(TIntrusivePtr> partitionStream, - std::vector::TMessage>& messages, - std::vector::TCompressedMessage>& compressedMessages, - size_t& maxByteSize, - size_t& dataSize) const +typename TDataDecompressionInfo::TDecompressedData +TDataDecompressionInfo::BuildDecompressedData(TIntrusivePtr> partitionStream, + size_t batchIndex, + size_t messageIndex, + const TDecompressionResult& codecResult) { - auto& msg = Parent->GetServerMessage(); + TDecompressedData result; + auto& msg = GetServerMessage(); i64 minOffset = Max(); i64 maxOffset = 0; - auto& batch = *msg.mutable_batches(Batch); - const auto& meta = Parent->GetBatchMeta(Batch); + auto& batch = *msg.mutable_batches(batchIndex); + const auto& meta = GetBatchMeta(batchIndex); const TInstant batchWriteTimestamp = [&batch](){ if constexpr (UseMigrationProtocol) { return TInstant::MilliSeconds(batch.write_timestamp_ms()); @@ -3053,40 +3129,78 @@ void TDataDecompressionEvent::TakeData(TIntrusivePtr(messageData.offset())); - maxOffset = Max(maxOffset, static_cast(messageData.offset())); + result.MessagesTaken = 1; - if constexpr (UseMigrationProtocol) { - using TMessageInformation = NPersQueue::TReadSessionEvent::TDataReceivedEvent::TMessageInformation; + if constexpr (!UseMigrationProtocol) { + const auto& messageMeta = GetMessageMeta(batchIndex, messageIndex); + + if (GetDoDecompress()) { + ui64 recordsSkipped = 0; + result.MessagesTaken = 0; + + for (const auto& decompressedMsg : codecResult.Messages) { + ui64 offset = messageData.offset(); + ui64 seqNo = messageData.seq_no(); + TInstant createTime = TInstant::MilliSeconds( + ::google::protobuf::util::TimeUtil::TimestampToMilliseconds(messageData.created_at())); + + if (decompressedMsg.Meta) { + const auto& recordMeta = *decompressedMsg.Meta; + offset = static_cast(messageData.offset()) + static_cast(recordMeta.OffsetDelta); + if (offset < CommittedOffset) { + ++result.MessagesTaken; + ++recordsSkipped; + continue; + } + seqNo = static_cast(*codecResult.BatchBaseSequence) + static_cast(recordMeta.SequenceDelta); + createTime = TInstant::MilliSeconds(*codecResult.BatchBaseTimestampMs + recordMeta.TimestampDelta); + } - TMessageInformation messageInfo(messageData.offset(), - batch.source_id(), - messageData.seq_no(), - TInstant::MilliSeconds(messageData.create_timestamp_ms()), - batchWriteTimestamp, - batch.ip(), - meta, - messageData.uncompressed_size()); + TReadSessionEvent::TDataReceivedEvent::TMessageInformation messageInfo( + offset, + batch.producer_id(), + seqNo, + createTime, + batchWriteTimestamp, + meta, + messageMeta, + decompressedMsg.Data.size(), + messageData.message_group_id() + ); + + minOffset = Min(minOffset, static_cast(offset)); + maxOffset = Max(maxOffset, static_cast(offset)); + result.DataSize += decompressedMsg.Data.size(); + + result.Messages.emplace_back( + std::move(decompressedMsg.Data), + GetDecompressionError(batchIndex, messageIndex), + std::move(messageInfo), + partitionStream); + + ++result.MessagesTaken; + } - if (Parent->GetDoDecompress()) { - messages.emplace_back(messageData.data(), - Parent->GetDecompressionError(Batch, Message), - messageInfo, - partitionStream, - messageData.partition_key(), - messageData.explicit_hash()); - } else { - compressedMessages.emplace_back(static_cast(messageData.codec()), - messageData.data(), - std::vector{messageInfo}, - partitionStream, - messageData.partition_key(), - messageData.explicit_hash()); + messageData.clear_data(); + + auto log = partitionStream->GetLog(); + if (recordsSkipped > 0) { + LOG_LAZY(log, TLOG_DEBUG, TStringBuilder() + << "Take Data (codec batch). Partition " << partitionStream->GetPartitionId() + << ". Read: {" << batchIndex << ", " << messageIndex << "} (" + << minOffset << "-" << maxOffset << "), messages: " << result.MessagesTaken + << ", skipped as committed: " << recordsSkipped); + } else { + LOG_LAZY(log, TLOG_DEBUG, TStringBuilder() + << "Take Data. Partition " << partitionStream->GetPartitionId() + << ". Read: {" << batchIndex << ", " << messageIndex << "} (" + << minOffset << "-" << maxOffset << ")"); + } + return result; } - } else { - const auto& messageMeta = Parent->GetMessageMeta(Batch, Message); + TReadSessionEvent::TDataReceivedEvent::TMessageInformation messageInfo( messageData.offset(), batch.producer_id(), @@ -3096,38 +3210,149 @@ void TDataDecompressionEvent::TakeData(TIntrusivePtrGetDoDecompress()) { - messages.emplace_back(messageData.data(), - Parent->GetDecompressionError(Batch, Message), - messageInfo, - partitionStream); - } else { - compressedMessages.emplace_back(static_cast(batch.codec()), - messageData.data(), - messageInfo, - partitionStream); - } + minOffset = Min(minOffset, static_cast(messageData.offset())); + maxOffset = Max(maxOffset, static_cast(messageData.offset())); + + result.CompressedMessages.emplace_back(static_cast(batch.codec()), + messageData.data(), + messageInfo, + partitionStream); + + result.DataSize += messageData.data().size(); + messageData.clear_data(); + + auto log = partitionStream->GetLog(); + LOG_LAZY(log, TLOG_DEBUG, TStringBuilder() + << "Take Data. Partition " << partitionStream->GetPartitionId() + << ". Read: {" << batchIndex << ", " << messageIndex << "} (" + << minOffset << "-" << maxOffset << ")"); + return result; } - maxByteSize -= Min(maxByteSize, messageData.data().size()); + minOffset = Min(minOffset, static_cast(messageData.offset())); + maxOffset = Max(maxOffset, static_cast(messageData.offset())); + + if constexpr (UseMigrationProtocol) { + using TMessageInformation = NPersQueue::TReadSessionEvent::TDataReceivedEvent::TMessageInformation; + + TMessageInformation messageInfo(messageData.offset(), + batch.source_id(), + messageData.seq_no(), + TInstant::MilliSeconds(messageData.create_timestamp_ms()), + batchWriteTimestamp, + batch.ip(), + meta, + messageData.uncompressed_size()); - dataSize += messageData.data().size(); + if (GetDoDecompress()) { + Y_ASSERT(!codecResult.Messages.empty()); + const auto& decompressedMsg = codecResult.Messages.front(); + result.Messages.emplace_back(decompressedMsg.Data, + GetDecompressionError(batchIndex, messageIndex), + messageInfo, + partitionStream, + messageData.partition_key(), + messageData.explicit_hash()); + result.DataSize += decompressedMsg.Data.size(); + } else { + result.CompressedMessages.emplace_back(static_cast(messageData.codec()), + messageData.data(), + std::vector{messageInfo}, + partitionStream, + messageData.partition_key(), + messageData.explicit_hash()); + result.DataSize += messageData.data().size(); + } + } // Clear data to free internal session's memory. messageData.clear_data(); - LOG_LAZY(partitionStream->GetLog(), TLOG_DEBUG, TStringBuilder() + auto log = partitionStream->GetLog(); + LOG_LAZY(log, TLOG_DEBUG, TStringBuilder() << "Take Data. Partition " << partitionStream->GetPartitionId() - << ". Read: {" << Batch << ", " << Message << "} (" + << ". Read: {" << batchIndex << ", " << messageIndex << "} (" << minOffset << "-" << maxOffset << ")"); + return result; +} + +template +size_t TDataDecompressionInfo::GetPreparedDataSize(size_t batch, size_t message) const +{ + Y_ASSERT(batch < DecompressedData.size()); + Y_ASSERT(message < DecompressedData[batch].size()); + const auto& data = DecompressedData[batch][message]; + if (data.MessagesTaken > 0) { + return data.DataSize; + } + return GetServerMessage().batches(batch).message_data(message).data().size(); +} + +template +size_t TDataDecompressionInfo::GetPreparedMessageCount(size_t batch, size_t message) const +{ + Y_ASSERT(batch < DecompressedData.size()); + Y_ASSERT(message < DecompressedData[batch].size()); + const auto& data = DecompressedData[batch][message]; + if (data.MessagesTaken > 0) { + return data.MessagesTaken; + } + if constexpr (UseMigrationProtocol) { + return 1; + } else { + const auto& serverBatch = GetServerMessage().batches(batch); + return GetReadMessageCount(serverBatch.message_data(message), serverBatch.codec()); + } +} + +template +void TDataDecompressionInfo::PutDecompressedData(size_t batch, + size_t message, + TDecompressedData&& data) +{ + Y_ASSERT(batch < DecompressedData.size()); + Y_ASSERT(message < DecompressedData[batch].size()); + DecompressedData[batch][message] = std::move(data); +} + +template +typename TDataDecompressionInfo::TDecompressedData +TDataDecompressionInfo::TakeData(TIntrusivePtr> partitionStream, + size_t batch, + size_t message, + size_t& maxByteSize) +{ + Y_ASSERT(batch < DecompressedData.size()); + Y_ASSERT(message < DecompressedData[batch].size()); + + TDecompressedData result = std::move(DecompressedData[batch][message]); + if (result.MessagesTaken == 0 && result.Messages.empty() && result.CompressedMessages.empty()) { + result = BuildDecompressedData(partitionStream, batch, message, TDecompressionResult{}); + } + maxByteSize -= Min(maxByteSize, result.DataSize); + return result; +} + +template +typename TDataDecompressionInfo::TDecompressedData +TDataDecompressionEvent::TakeData(TIntrusivePtr> partitionStream, + size_t& maxByteSize) const +{ + return Parent->TakeData(partitionStream, Batch, Message, maxByteSize); } template size_t TDataDecompressionEvent::GetDataSize() const { - return Parent->GetServerMessage().batches(Batch).message_data(Message).data().size(); + return Parent->GetPreparedDataSize(Batch, Message); +} + +template +size_t TDataDecompressionEvent::GetMessageCount() const { + return Parent->GetPreparedMessageCount(Batch, Message); } template @@ -3187,6 +3412,7 @@ TDataDecompressionInfo::TDecompressionTask::TDecompression : Parent(std::move(parent)) , PartitionStream(std::move(partitionStream)) , Ready(ready) { + Y_ABORT_UNLESS(PartitionStream); } template @@ -3204,30 +3430,39 @@ void TDataDecompressionInfo::TDecompressionTask::operator( for (size_t i = messages.MessageRange.first; i < messages.MessageRange.second; ++i) { auto& data = *batch.mutable_message_data(i); - ++messagesProcessed; dataProcessed += static_cast(data.data().size()); minOffset = Min(minOffset, static_cast(data.offset())); maxOffset = Max(maxOffset, static_cast(data.offset())); + TDecompressionResult codecResult; try { - if constexpr (UseMigrationProtocol) { - if (parent->DoDecompress - && data.codec() != Ydb::PersQueue::V1::CODEC_RAW - && data.codec() != Ydb::PersQueue::V1::CODEC_UNSPECIFIED - ) { - const ICodec* codecImpl = TCodecMap::GetTheCodecMap().GetOrThrow(static_cast(data.codec())); - std::string decompressed = codecImpl->Decompress(data.data()); - data.set_data(TStringType{std::move(decompressed)}); - data.set_codec(Ydb::PersQueue::V1::CODEC_RAW); - } - } else { - if (parent->DoDecompress - && static_cast(batch.codec()) != Ydb::Topic::CODEC_RAW - && static_cast(batch.codec()) != Ydb::Topic::CODEC_UNSPECIFIED - ) { - const ICodec* codecImpl = TCodecMap::GetTheCodecMap().GetOrThrow(static_cast(batch.codec())); - std::string decompressed = codecImpl->Decompress(data.data()); - data.set_data(TStringType{std::move(decompressed)}); + if (parent->DoDecompress) { + if constexpr (UseMigrationProtocol) { + if (data.codec() != Ydb::PersQueue::V1::CODEC_RAW + && data.codec() != Ydb::PersQueue::V1::CODEC_UNSPECIFIED + ) { + const ICodec* codecImpl = TCodecMap::GetTheCodecMap().GetOrThrow(static_cast(data.codec())); + codecResult = codecImpl->DecompressData(data.data()); + data.set_codec(Ydb::PersQueue::V1::CODEC_RAW); + } else { + codecResult.Messages.push_back(TDecompressedMessage{ + .Data = std::string(data.data()), + .Meta = std::nullopt, + }); + } + } else { + const auto batchCodec = static_cast(batch.codec()); + if (batchCodec != Ydb::Topic::CODEC_RAW + && batchCodec != Ydb::Topic::CODEC_UNSPECIFIED + ) { + const ICodec* codecImpl = TCodecMap::GetTheCodecMap().GetOrThrow(static_cast(batch.codec())); + codecResult = codecImpl->DecompressData(data.data()); + } else { + codecResult.Messages.push_back(TDecompressedMessage{ + .Data = std::string(data.data()), + .Meta = std::nullopt, + }); + } } } } catch (...) { @@ -3236,13 +3471,24 @@ void TDataDecompressionInfo::TDecompressionTask::operator( if (auto session = parent->CbContext->LockShared()) { session->GetLog() << TLOG_INFO << "Error decompressing data: " << CurrentExceptionMessage(); } + + if (codecResult.Messages.empty()) { + codecResult.Messages.push_back(TDecompressedMessage{ + .Data = std::string(data.data()), + .Meta = std::nullopt, + }); + } } - DecompressedSize += data.data().size(); + auto decompressedData = parent->BuildDecompressedData(PartitionStream, messages.Batch, i, codecResult); + DecompressedSize += decompressedData.DataSize; + messagesProcessed += decompressedData.MessagesTaken; + parent->PutDecompressedData(messages.Batch, i, std::move(decompressedData)); } } if (auto session = parent->CbContext->LockShared()) { + const auto& log = session->GetLog(); const i64 partition_id = [parent](){ if constexpr (UseMigrationProtocol) { return parent->ServerMessage.partition(); @@ -3250,9 +3496,9 @@ void TDataDecompressionInfo::TDecompressionTask::operator( return parent->ServerMessage.partition_session_id(); } }(); - LOG_LAZY(session->GetLog(), TLOG_DEBUG, TStringBuilder() << "Decompression task done. Partition/PartitionSessionId: " - << partition_id << " (" << minOffset << "-" - << maxOffset << ")"); + LOG_LAZY(log, TLOG_DEBUG, TStringBuilder() << "Decompression task done. Partition/PartitionSessionId: " + << partition_id << " (" << minOffset << "-" + << maxOffset << ")"); } Y_ASSERT(dataProcessed == SourceDataSize); @@ -3282,12 +3528,15 @@ void TDataDecompressionInfo::TDecompressionTask::ClearPare // TUserRetrievedEventsInfoAccumulator template -void TUserRetrievedEventsInfoAccumulator::Add(TDataDecompressionInfoPtr info, i64 decompressedSize) +void TUserRetrievedEventsInfoAccumulator::Add( + TDataDecompressionInfoPtr info, + i64 decompressedSize, + size_t messagesCount) { auto& counter = Counters[info]; counter.DecompressedSize += decompressedSize; - ++counter.MessagesCount; + counter.MessagesCount += messagesCount; } template diff --git a/src/client/topic/impl/write_session_impl.cpp b/src/client/topic/impl/write_session_impl.cpp index b8ebde1964a..0e737dea2a7 100644 --- a/src/client/topic/impl/write_session_impl.cpp +++ b/src/client/topic/impl/write_session_impl.cpp @@ -2,6 +2,7 @@ #include #include +#include #include @@ -9,6 +10,7 @@ #include #include +#include #include #include @@ -27,6 +29,25 @@ const uint64_t WRITE_ERROR_PARTITION_INACTIVE = 500029; namespace { using TTxId = std::pair; + +bool ValidateWriteSessionSettings(const TWriteSessionSettings& settings, NYdb::NIssue::TIssues& issues) { + if (!settings.BatchInnerCodec_.has_value()) { + return true; + } + + if (settings.Codec_ != ECodec::KAFKA_BATCH) { + issues.AddIssue("BatchInnerCodec can be set only when Codec is KAFKA_BATCH."); + return false; + } + + const ECodec innerCodec = *settings.BatchInnerCodec_; + if (innerCodec != ECodec::GZIP && innerCodec != ECodec::ZSTD) { + issues.AddIssue("BatchInnerCodec supports only GZIP and ZSTD when Codec is KAFKA_BATCH."); + return false; + } + + return true; +} using TTxIdOpt = std::optional; TTxIdOpt GetTransactionId(const Ydb::Topic::StreamWriteMessage_WriteRequest& request) @@ -82,10 +103,11 @@ TWriteSessionImpl::TWriteSessionImpl( , Connections(std::move(connections)) , DbDriverState(std::move(dbDriverState)) , PrevToken(DbDriverState->CredentialsProvider ? DbDriverState->CredentialsProvider->GetAuthInfo() : "") + , MaxBlockMessageCount(Settings.BatchFlushMessageCount_) , InitSeqNoPromise(NThreading::NewPromise()) , WakeupInterval( - Settings.BatchFlushInterval_.value_or(TDuration::Zero()) ? - std::min(Settings.BatchFlushInterval_.value_or(TDuration::Seconds(1)) / 5, TDuration::MilliSeconds(100)) + Settings.BatchFlushInterval_ != TDuration::Zero() ? + std::min(Settings.BatchFlushInterval_ / 5, TDuration::MilliSeconds(100)) : TDuration::MilliSeconds(100) ) @@ -122,6 +144,16 @@ void TWriteSessionImpl::Start(const TDuration& delay) { } if (!Started) { + NYdb::NIssue::TIssues issues; + if (!ValidateWriteSessionSettings(Settings, issues)) { + with_lock(Lock) { + CloseImpl( + EStatus::BAD_REQUEST, + MakeIssueWithSubIssues("Invalid write session settings", issues)); + } + return; + } + with_lock(Lock) { HandleWakeUpImpl(); } @@ -657,8 +689,8 @@ void TWriteSessionImpl::WriteInternal(TContinuationToken&&, TWriteMessage&& mess MakeTransactionId(message.GetTxPtr()) ); - FlushWriteIfRequiredImpl(); readyToAccept = OnMemoryUsageChangedImpl(static_cast(bufferSize)).NowOk; + FlushWriteIfRequiredImpl(); } if (readyToAccept) { EventsQueue->PushEvent(TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()}); @@ -873,6 +905,7 @@ void TWriteSessionImpl::InitImpl() { for (const auto& attr : Settings.Meta_.Fields) { (*init->mutable_write_session_meta())[attr.first] = attr.second; } + LOG_LAZY(DbDriverState->Log, TLOG_DEBUG, LogPrefixImpl() << "Write session: send init request: "<< req.ShortDebugString()); TRACE_LAZY(DbDriverState->Log, "InitRequest", @@ -1064,8 +1097,10 @@ TWriteSessionImpl::TProcessSrvMessageResult TWriteSessionImpl::ProcessServerMess LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefixImpl() << "Write session established. Init response: " << initResponse.ShortDebugString()); TRACE_LAZY(DbDriverState->Log, "InitResponse", TRACE_KV("partition_id", initResponse.partition_id()), - TRACE_KV("session_id", initResponse.session_id())); + TRACE_KV("session_id", initResponse.session_id()), + TRACE_KV("is_batching_supported", initResponse.is_batching_supported())); SessionId = initResponse.session_id(); + BatchingSupported = initResponse.is_batching_supported(); auto prevDirectWriteToPartitionId = DirectWriteToPartitionId; if (Settings.DirectWriteToPartition_ && !Settings.PartitionId_.has_value()) { @@ -1124,12 +1159,14 @@ TWriteSessionImpl::TProcessSrvMessageResult TWriteSessionImpl::ProcessServerMess writeStat->TopicQuotedTime = durationConv(stat.topic_quota_wait_time()); for (const auto& ack : batchWriteResponse.acks()) { - // TODO: Fill writer statistics - uint64_t msgId = GetIdImpl(ack.seq_no()); + const uint64_t firstMsgId = GetIdImpl(ack.seq_no()); - Y_ABORT_UNLESS(ack.has_written() || ack.has_skipped() || ack.has_written_in_tx()); + size_t ackCount = 1; + if (!SentPackedMessage.empty() && SentPackedMessage.front().Offset == firstMsgId) { + ackCount = SentPackedMessage.front().MessageCount; + } - TrySignalAllAcksReceived(msgId); + Y_ABORT_UNLESS(ack.has_written() || ack.has_skipped() || ack.has_written_in_tx()); TWriteSessionEvent::TWriteAck::EEventState msgWriteStatus; if (ack.has_written_in_tx()) { @@ -1143,19 +1180,24 @@ TWriteSessionImpl::TProcessSrvMessageResult TWriteSessionImpl::ProcessServerMess : TWriteSessionEvent::TWriteAck::EES_DISCARDED; } - uint64_t offset = ack.has_written() ? ack.written().offset() : 0; + const uint64_t baseOffset = ack.has_written() ? ack.written().offset() : 0; - acksEvent.Acks.push_back(TWriteSessionEvent::TWriteAck{ - msgId, - msgWriteStatus, - TWriteSessionEvent::TWriteAck::TWrittenMessageDetails { - offset, - PartitionId, - }, - writeStat, - }); + for (size_t i = 0; i < ackCount; ++i) { + const uint64_t msgId = firstMsgId + i; + TrySignalAllAcksReceived(msgId); + + acksEvent.Acks.push_back(TWriteSessionEvent::TWriteAck{ + msgId, + msgWriteStatus, + TWriteSessionEvent::TWriteAck::TWrittenMessageDetails { + baseOffset + i, + PartitionId, + }, + writeStat, + }); + } - if (CleanupOnAcknowledgedImpl(msgId)) { + if (CleanupOnAcknowledgedImpl(firstMsgId)) { result.Events.emplace_back(TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()}); } } @@ -1182,14 +1224,17 @@ bool TWriteSessionImpl::CleanupOnAcknowledgedImpl(uint64_t id) { const auto& sentFront = SentOriginalMessages.front(); uint64_t size = 0; uint64_t compressedSize = 0; + size_t packedMessageCount = 1; if(!SentPackedMessage.empty() && SentPackedMessage.front().Offset == id) { - auto memoryUsage = OnMemoryUsageChangedImpl(-static_cast(SentPackedMessage.front().Data.size())); - result = memoryUsage.NowOk && !memoryUsage.WasOk; const auto& front = SentPackedMessage.front(); + const ui64 memoryToRelease = front.Compressed ? front.Data.size() : front.OriginalMemoryUsage; + auto memoryUsage = OnMemoryUsageChangedImpl(-static_cast(memoryToRelease)); + result = memoryUsage.NowOk && !memoryUsage.WasOk; + packedMessageCount = front.MessageCount; if (front.Compressed) { compressedSize = front.Data.size(); } else { - size = front.Data.size(); + size = front.OriginalSize; } (*Counters->MessagesWritten) += front.MessageCount; @@ -1211,12 +1256,20 @@ bool TWriteSessionImpl::CleanupOnAcknowledgedImpl(uint64_t id) { Y_ABORT_UNLESS(Counters->BytesInflightCompressed->Val() >= 0); Y_ABORT_UNLESS(Counters->BytesInflightUncompressed->Val() >= 0); - Y_ABORT_UNLESS(sentFront.Id == id); + if (packedMessageCount > 1) { + for (size_t i = 0; i < packedMessageCount; ++i) { + Y_ABORT_UNLESS(!SentOriginalMessages.empty()); + Y_ABORT_UNLESS(SentOriginalMessages.front().Id == id + i); + WrittenInTx.erase(SentOriginalMessages.front().Id); + SentOriginalMessages.pop(); + } + } else { + Y_ABORT_UNLESS(sentFront.Id == id); + WrittenInTx.erase(id); + SentOriginalMessages.pop(); + } (*Counters->BytesInflightTotal) = MemoryUsage; - SentOriginalMessages.pop(); - - WrittenInTx.erase(id); return result; } @@ -1275,23 +1328,32 @@ void TWriteSessionImpl::CompressImpl(TBlock&& block_) { } Y_ABORT_UNLESS(block_.Valid); + const i64 baseSequence = static_cast(GetSeqNoImpl(block_.Offset)); std::shared_ptr blockPtr(std::make_shared()); blockPtr->Move(block_); auto lambda = [cbContext = SelfContext, codec = Settings.Codec_, + batchInnerCodec = Settings.BatchInnerCodec_, level = Settings.CompressionLevel_, + baseSequence, isSyncCompression = !CompressionExecutor->IsAsync(), blockPtr, client = Client]() mutable { Y_ABORT_UNLESS(!blockPtr->Compressed); - auto compressedData = CompressBuffer( - std::move(client), blockPtr->OriginalDataRefs, codec, level - ); - Y_ABORT_UNLESS(!compressedData.Empty()); - blockPtr->Data = std::move(compressedData); - blockPtr->Compressed = true; - blockPtr->CodecID = static_cast(codec); + const ICodec* codecImpl = TCodecMap::GetTheCodecMap().GetOrThrow(static_cast(codec)); + TWriteBlockCompression compression{ + .Codec = codec, + .Payloads = blockPtr->OriginalDataRefs, + .CreatedAt = blockPtr->CreatedAt, + .Data = blockPtr->Data, + .CodecID = blockPtr->CodecID, + .Compressed = blockPtr->Compressed, + .BaseSequence = baseSequence, + .BatchInnerCodec = batchInnerCodec, + .CompressionLevel = level, + }; + codecImpl->CompressWriteBlock(compression); if (auto self = cbContext->LockShared()) { self->OnCompressed(std::move(*blockPtr), isSyncCompression); } @@ -1318,9 +1380,13 @@ TMemoryUsageChange TWriteSessionImpl::OnCompressedImpl(TBlock&& block) { UpdateTimedCountersImpl(); Y_ABORT_UNLESS(block.Valid); - auto memoryUsage = OnMemoryUsageChangedImpl(static_cast(block.Data.size()) - static_cast(block.OriginalMemoryUsage)); - (*Counters->BytesInflightUncompressed) -= block.OriginalSize; - (*Counters->BytesInflightCompressed) += block.Data.size(); + + TMemoryUsageChange memoryUsage{MemoryUsage <= Settings.MaxMemoryUsage_, MemoryUsage <= Settings.MaxMemoryUsage_}; + if (block.Compressed) { + memoryUsage = OnMemoryUsageChangedImpl(static_cast(block.Data.size()) - static_cast(block.OriginalMemoryUsage)); + (*Counters->BytesInflightUncompressed) -= block.OriginalSize; + (*Counters->BytesInflightCompressed) += block.Data.size(); + } PackedMessagesToSend.emplace(std::move(block)); @@ -1349,15 +1415,15 @@ void TWriteSessionImpl::ResetForRetryImpl() { SentPackedMessage.pop(); } uint64_t minId = PackedMessagesToSend.empty() ? NextId + 1 : PackedMessagesToSend.top().Offset; - std::queue freshOriginalMessagesToSend; + std::deque freshOriginalMessagesToSend; OriginalMessagesToSend.swap(freshOriginalMessagesToSend); while (!SentOriginalMessages.empty()) { - OriginalMessagesToSend.emplace(std::move(SentOriginalMessages.front())); + OriginalMessagesToSend.emplace_back(std::move(SentOriginalMessages.front())); SentOriginalMessages.pop(); } while (!freshOriginalMessagesToSend.empty()) { - OriginalMessagesToSend.emplace(std::move(freshOriginalMessagesToSend.front())); - freshOriginalMessagesToSend.pop(); + OriginalMessagesToSend.emplace_back(std::move(freshOriginalMessagesToSend.front())); + freshOriginalMessagesToSend.pop_front(); } if (!OriginalMessagesToSend.empty() && OriginalMessagesToSend.front().Id < minId) minId = OriginalMessagesToSend.front().Id; @@ -1371,8 +1437,10 @@ void TWriteSessionImpl::FlushWriteIfRequiredImpl() { if (!CurrentBatch.Empty() && !CurrentBatch.FlushRequested) { MessagesAcquired += static_cast(CurrentBatch.Acquire()); - if (TInstant::Now() - CurrentBatch.StartedAt >= Settings.BatchFlushInterval_.value_or(TDuration::Zero()) - || CurrentBatch.CurrentSize >= Settings.BatchFlushSizeBytes_.value_or(0) + if (TInstant::Now() - CurrentBatch.StartedAt >= Settings.BatchFlushInterval_ + || (Settings.BatchFlushSizeBytes_.has_value() + && (Settings.BatchFlushSizeBytes_.value() == 0 + || CurrentBatch.CurrentSize >= Settings.BatchFlushSizeBytes_.value())) || CurrentBatch.CurrentSize >= MaxBlockSize || CurrentBatch.Messages.size() >= MaxBlockMessageCount || CurrentBatch.HasCodec() @@ -1394,6 +1462,24 @@ size_t TWriteSessionImpl::WriteBatchImpl() { Y_ABORT_UNLESS(CurrentBatch.Messages.size() <= MaxBlockMessageCount); + while (!CurrentBatch.Acquired) { + if (!CurrentBatch.Acquire()) { + break; + } + } + + // TBuffer::Append in Acquire() may reallocate the batch buffer, invalidating + // DataRefs of previously acquired messages. Payloads are laid out in the buffer + // contiguously in message order, so re-base all refs onto the final buffer. + { + size_t dataOffset = 0; + for (auto& message : CurrentBatch.Messages) { + Y_ABORT_UNLESS(dataOffset + message.DataRef.size() <= CurrentBatch.Data.size()); + message.DataRef = std::string_view(CurrentBatch.Data.data() + dataOffset, message.DataRef.size()); + dataOffset += message.DataRef.size(); + } + } + const bool skipCompression = Settings.Codec_ == ECodec::RAW || CurrentBatch.HasCodec(); if (!skipCompression && Settings.CompressionExecutor_->IsAsync()) { MessagesAcquired += static_cast(CurrentBatch.Acquire()); @@ -1419,28 +1505,35 @@ size_t TWriteSessionImpl::WriteBatchImpl() { block.MessageCount += 1; const auto& datum = currMessage.DataRef; block.OriginalSize += datum.size(); - block.OriginalMemoryUsage = CurrentBatch.Data.size(); + block.OriginalMemoryUsage += datum.size(); block.OriginalDataRefs.emplace_back(datum); + block.CreatedAt.emplace_back(createTs); if (CurrentBatch.Messages[i].Codec.has_value()) { Y_ABORT_UNLESS(CurrentBatch.Messages.size() == 1); block.CodecID = static_cast(*currMessage.Codec); block.OriginalSize = currMessage.OriginalSize; - block.Compressed = false; + block.Compressed = true; } size += datum.size(); UpdateTimedCountersImpl(); - (*Counters->BytesInflightUncompressed) += datum.size(); + if (block.Compressed) { + (*Counters->BytesInflightCompressed) += datum.size(); + } else { + (*Counters->BytesInflightUncompressed) += datum.size(); + } (*Counters->MessagesInflight)++; if (!currMessage.MessageMeta.empty()) { - OriginalMessagesToSend.emplace(id, createTs, datum.size(), + OriginalMessagesToSend.emplace_back(id, createTs, datum.size(), std::move(currMessage.MessageMeta), std::move(currMessage.Tx)); } else { - OriginalMessagesToSend.emplace(id, createTs, datum.size(), + OriginalMessagesToSend.emplace_back(id, createTs, datum.size(), std::move(currMessage.Tx)); } } + block.Data = std::move(CurrentBatch.Data); + if (skipCompression) { PackedMessagesToSend.emplace(std::move(block)); } else { @@ -1454,9 +1547,7 @@ size_t TWriteSessionImpl::WriteBatchImpl() { return size; } -size_t GetMaxGrpcMessageSize() { - return 120_MB; -} +namespace NGrpc = NWriteSessionGrpc; bool TWriteSessionImpl::IsReadyToSendNextImpl() const { Y_ABORT_UNLESS(Lock.IsLocked()); @@ -1513,6 +1604,105 @@ bool TWriteSessionImpl::TxIsChanged(const Ydb::Topic::StreamWriteMessage_WriteRe return GetTransactionId(*writeRequest) != GetTransactionId(OriginalMessagesToSend.front().Tx); } +void TWriteSessionImpl::SendBatchBlock( + const TBlock& block, + Ydb::Topic::StreamWriteMessage_WriteRequest* writeRequest) +{ + Y_ABORT_UNLESS(Lock.IsLocked()); + Y_ABORT_UNLESS(writeRequest); + Y_ABORT_UNLESS(block.MessageCount > 1); + + if (!BatchingSupported) { + ThrowFatalError("Server does not support messages batching"); + } + + Y_ABORT_UNLESS(!OriginalMessagesToSend.empty()); + std::vector batchMessages; + batchMessages.reserve(block.MessageCount); + for (size_t i = 0; i < block.MessageCount; ++i) { + Y_ABORT_UNLESS(!OriginalMessagesToSend.empty()); + batchMessages.emplace_back(std::move(OriginalMessagesToSend.front())); + OriginalMessagesToSend.pop_front(); + } + + const auto& firstMessage = batchMessages.front(); + auto* msgData = writeRequest->add_messages(); + + if (firstMessage.Tx) { + writeRequest->mutable_tx()->set_id(firstMessage.Tx->TxId); + writeRequest->mutable_tx()->set_session(firstMessage.Tx->SessionId); + } + + msgData->set_seq_no(static_cast(GetSeqNoImpl(batchMessages.back().Id))); + *msgData->mutable_created_at() = + ::google::protobuf::util::TimeUtil::MillisecondsToTimestamp(firstMessage.CreatedAt.MilliSeconds()); + + for (auto& [k, v] : firstMessage.MessageMeta) { + auto* pair = msgData->add_metadata_items(); + pair->set_key(TStringType{k}); + pair->set_value(TStringType{v}); + } + for (size_t i = 1; i < batchMessages.size(); ++i) { + for (auto& [k, v] : batchMessages[i].MessageMeta) { + if (k != NGrpc::PARTITION_KEY_META_KEY) { + continue; + } + auto* pair = msgData->add_metadata_items(); + pair->set_key(TStringType{k}); + pair->set_value(TStringType{v}); + } + } + + msgData->set_uncompressed_size(static_cast(block.OriginalSize)); + msgData->set_data(block.Data.data(), block.Data.size()); + + for (auto& message : batchMessages) { + SentOriginalMessages.emplace(std::move(message)); + } +} + +void TWriteSessionImpl::SendStandardBlock( + const TBlock& block, + Ydb::Topic::StreamWriteMessage_WriteRequest* writeRequest) +{ + Y_ABORT_UNLESS(Lock.IsLocked()); + Y_ABORT_UNLESS(writeRequest); + Y_ABORT_UNLESS(block.MessageCount == 1); + + for (size_t i = 0; i != block.MessageCount; ++i) { + Y_ABORT_UNLESS(!OriginalMessagesToSend.empty()); + + auto& message = OriginalMessagesToSend.front(); + auto* msgData = writeRequest->add_messages(); + + if (message.Tx) { + writeRequest->mutable_tx()->set_id(message.Tx->TxId); + writeRequest->mutable_tx()->set_session(message.Tx->SessionId); + } + + msgData->set_seq_no(GetSeqNoImpl(message.Id)); + *msgData->mutable_created_at() = + ::google::protobuf::util::TimeUtil::MillisecondsToTimestamp(message.CreatedAt.MilliSeconds()); + + for (auto& [k, v] : message.MessageMeta) { + auto* pair = msgData->add_metadata_items(); + pair->set_key(TStringType{k}); + pair->set_value(TStringType{v}); + } + SentOriginalMessages.emplace(std::move(message)); + OriginalMessagesToSend.pop_front(); + + msgData->set_uncompressed_size(block.OriginalSize); + if (block.Compressed) { + msgData->set_data(block.Data.data(), block.Data.size()); + } else { + for (auto& buffer : block.OriginalDataRefs) { + msgData->set_data(buffer.data(), buffer.size()); + } + } + } +} + void TWriteSessionImpl::SendImpl() { Y_ABORT_UNLESS(Lock.IsLocked()); @@ -1523,10 +1713,10 @@ void TWriteSessionImpl::SendImpl() { ui32 prevCodec = 0; - ui64 currentSize = 0; + NGrpc::TRequestSizeLimiter sizeLimiter(2); // Send blocks while we can without messages reordering. - while (IsReadyToSendNextImpl() && currentSize < GetMaxGrpcMessageSize()) { + while (IsReadyToSendNextImpl()) { const auto& block = PackedMessagesToSend.top(); Y_ABORT_UNLESS(block.Valid); if (writeRequest->messages_size() > 0 && prevCodec != block.CodecID) { @@ -1535,47 +1725,28 @@ void TWriteSessionImpl::SendImpl() { if (TxIsChanged(writeRequest)) { break; } - prevCodec = block.CodecID; - writeRequest->set_codec(static_cast(block.CodecID)); - Y_ABORT_UNLESS(block.MessageCount == 1); - for (size_t i = 0; i != block.MessageCount; ++i) { - Y_ABORT_UNLESS(!OriginalMessagesToSend.empty()); - - auto& message = OriginalMessagesToSend.front(); - auto* msgData = writeRequest->add_messages(); - if (message.Tx) { - writeRequest->mutable_tx()->set_id(message.Tx->TxId); - writeRequest->mutable_tx()->set_session(message.Tx->SessionId); - } + const size_t blockSize = NGrpc::EstimateTopicWriteRequestBlockSize(block, OriginalMessagesToSend, sizeLimiter.Empty()); + if (!sizeLimiter.CanAdd(blockSize)) { + break; + } - msgData->set_seq_no(GetSeqNoImpl(message.Id)); - *msgData->mutable_created_at() = ::google::protobuf::util::TimeUtil::MillisecondsToTimestamp(message.CreatedAt.MilliSeconds()); + prevCodec = block.CodecID; + writeRequest->set_codec(static_cast(block.CodecID)); - for (auto& [k, v] : message.MessageMeta) { - auto* pair = msgData->add_metadata_items(); - pair->set_key(TStringType{k}); - pair->set_value(TStringType{v}); - } - SentOriginalMessages.emplace(std::move(message)); - OriginalMessagesToSend.pop(); + const bool sendAsBatchBlock = block.MessageCount > 1; - msgData->set_uncompressed_size(block.OriginalSize); - if (block.Compressed) { - msgData->set_data(block.Data.data(), block.Data.size()); - } else { - for (auto& buffer: block.OriginalDataRefs) { - msgData->set_data(buffer.data(), buffer.size()); - } - } + if (sendAsBatchBlock) { + SendBatchBlock(block, writeRequest); + } else { + SendStandardBlock(block, writeRequest); } TBlock moveBlock; moveBlock.Move(block); SentPackedMessage.emplace(std::move(moveBlock)); PackedMessagesToSend.pop(); - - currentSize += writeRequest->ByteSizeLong(); + sizeLimiter.Add(blockSize); } UpdateTokenIfNeededImpl(); LOG_LAZY(DbDriverState->Log, @@ -1596,6 +1767,12 @@ bool TWriteSessionImpl::Close(TDuration closeTimeout) { { std::lock_guard guard(Lock); LOG_LAZY(DbDriverState->Log, TLOG_INFO, LogPrefixImpl() << "Write session: close. Timeout " << closeTimeout); + // Messages buffered in CurrentBatch are not yet in OriginalMessagesToSend, + // so the wait loop below would not see them. Flush them to avoid losing + // a partially filled batch on close. + if (!CurrentBatch.Empty()) { + WriteBatchImpl(); + } } auto startTime = TInstant::Now(); auto remaining = closeTimeout; diff --git a/src/client/topic/impl/write_session_impl.h b/src/client/topic/impl/write_session_impl.h index 0434ed07e67..313352ee441 100644 --- a/src/client/topic/impl/write_session_impl.h +++ b/src/client/topic/impl/write_session_impl.h @@ -233,6 +233,7 @@ class TWriteSessionImpl : public TContinuationTokenIssuer, size_t OriginalMemoryUsage = 0; ui32 CodecID = static_cast(ECodec::RAW); mutable std::vector OriginalDataRefs; + mutable std::vector CreatedAt; mutable TBuffer Data; bool Compressed = false; mutable bool Valid = true; @@ -250,11 +251,13 @@ class TWriteSessionImpl : public TContinuationTokenIssuer, OriginalMemoryUsage = rhs.OriginalMemoryUsage; CodecID = rhs.CodecID; OriginalDataRefs.swap(rhs.OriginalDataRefs); + CreatedAt.swap(rhs.CreatedAt); Data.Swap(rhs.Data); Compressed = rhs.Compressed; rhs.Data.Clear(); rhs.OriginalDataRefs.clear(); + rhs.CreatedAt.clear(); } }; @@ -405,6 +408,8 @@ class TWriteSessionImpl : public TContinuationTokenIssuer, uint64_t GetSeqNoImpl(uint64_t id); uint64_t GetIdImpl(uint64_t seqNo); void SendImpl(); + void SendBatchBlock(const TBlock& block, Ydb::Topic::StreamWriteMessage_WriteRequest* writeRequest); + void SendStandardBlock(const TBlock& block, Ydb::Topic::StreamWriteMessage_WriteRequest* writeRequest); void AbortImpl(); void CloseImpl(EStatus statusCode, NYdb::NIssue::TIssues&& issues); void CloseImpl(EStatus statusCode, const std::string& message); @@ -461,19 +466,20 @@ class TWriteSessionImpl : public TContinuationTokenIssuer, TMessageBatch CurrentBatch; - std::queue OriginalMessagesToSend; + std::deque OriginalMessagesToSend; std::priority_queue, Greater> PackedMessagesToSend; //! Messages that are sent but yet not acknowledged std::queue SentOriginalMessages; std::queue SentPackedMessage; const size_t MaxBlockSize = std::numeric_limits::max(); - const size_t MaxBlockMessageCount = 1; //!< Max message count that can be packed into a single block. In block version 0 is equal to 1 for compatibility + const size_t MaxBlockMessageCount; bool Connected = false; bool Started = false; std::atomic SendImplScheduled = false; std::atomic Aborting = 0; bool SessionEstablished = false; + bool BatchingSupported = false; ui32 PartitionId = 0; TPartitionLocation PreferredPartitionLocation = {}; uint64_t NextId = 0; diff --git a/src/client/topic/ut/write_session_size_ut.cpp b/src/client/topic/ut/write_session_size_ut.cpp new file mode 100644 index 00000000000..7391b53dbc7 --- /dev/null +++ b/src/client/topic/ut/write_session_size_ut.cpp @@ -0,0 +1,171 @@ +#include + +#include + +#include + +#include + +namespace NYdb::inline V3::NTopic::NTests { +namespace { + +namespace NGrpc = NWriteSessionGrpc; + +struct TTestBlock { + size_t MessageCount = 0; + ui32 CodecID = static_cast(ECodec::RAW); + size_t OriginalSize = 0; + std::vector OriginalDataRefs; + std::string Data; + bool Compressed = false; +}; + +struct TTestMessage { + TInstant CreatedAt; + std::vector> MessageMeta; + std::optional Tx; +}; + +void FillMetadata( + Ydb::Topic::StreamWriteMessage::WriteRequest::MessageData& message, + const std::vector>& metadata) { + for (const auto& [key, value] : metadata) { + auto* item = message.add_metadata_items(); + item->set_key(key); + item->set_value(value); + } +} + +void FillTx(Ydb::Topic::StreamWriteMessage::WriteRequest& request, const std::optional& tx) { + if (!tx) { + return; + } + request.mutable_tx()->set_id(tx->TxId); + request.mutable_tx()->set_session(tx->SessionId); +} + +size_t BuildTopicWriteRequestBlockSize( + const TTestBlock& block, + const std::vector& messages, + bool includeRequestFields) { + Y_ABORT_UNLESS(!messages.empty()); + + Ydb::Topic::StreamWriteMessage::WriteRequest request; + if (includeRequestFields) { + request.set_codec(static_cast(block.CodecID)); + FillTx(request, messages.front().Tx); + } + + if (block.MessageCount > 1) { + auto* message = request.add_messages(); + message->set_seq_no(-1); + *message->mutable_created_at() = + google::protobuf::util::TimeUtil::MillisecondsToTimestamp(messages.front().CreatedAt.MilliSeconds()); + FillMetadata(*message, messages.front().MessageMeta); + for (size_t i = 1; i < block.MessageCount; ++i) { + for (const auto& item : messages[i].MessageMeta) { + if (item.first == NGrpc::PARTITION_KEY_META_KEY) { + FillMetadata(*message, {item}); + } + } + } + message->set_uncompressed_size(static_cast(block.OriginalSize)); + message->set_data(block.Data); + return request.ByteSizeLong(); + } + + auto* message = request.add_messages(); + message->set_seq_no(-1); + *message->mutable_created_at() = + google::protobuf::util::TimeUtil::MillisecondsToTimestamp(messages.front().CreatedAt.MilliSeconds()); + FillMetadata(*message, messages.front().MessageMeta); + message->set_uncompressed_size(static_cast(block.OriginalSize)); + if (block.Compressed) { + message->set_data(block.Data); + } else { + Y_ABORT_UNLESS(block.OriginalDataRefs.size() == 1); + message->set_data(block.OriginalDataRefs.front().data(), block.OriginalDataRefs.front().size()); + } + return request.ByteSizeLong(); +} + +void AssertEstimateCloseToActual(size_t estimate, size_t actual) { + static constexpr size_t MaxAllowedOverestimate = 16; + + UNIT_ASSERT_GE_C(estimate, actual, "estimate=" << estimate << ", actual=" << actual); + UNIT_ASSERT_LE_C(estimate, actual + MaxAllowedOverestimate, "estimate=" << estimate << ", actual=" << actual); +} + +} // namespace + +Y_UNIT_TEST_SUITE(WriteSessionGrpcSize) { + Y_UNIT_TEST(RequestSizeLimiterBoundaries) { + const size_t maxSize = NGrpc::ProtoMessageFieldSize(2, 10); + NGrpc::TRequestSizeLimiter limiter(2, maxSize); + + // The first block is allowed even when it is larger than the limit: + // otherwise a single oversized block would make the send loop stuck. + UNIT_ASSERT(limiter.Empty()); + UNIT_ASSERT(limiter.CanAdd(maxSize * 2)); + + limiter.Add(5); + UNIT_ASSERT(!limiter.Empty()); + + // For non-empty requests the limiter checks the full serialized + // envelope size: field #2 tag + body length + body bytes. + UNIT_ASSERT(limiter.CanAdd(5)); + UNIT_ASSERT(!limiter.CanAdd(6)); + } + + Y_UNIT_TEST(EstimateTopicWriteRequestStandardBlockSize) { + const std::string data = "payload"; + TTestBlock block; + block.MessageCount = 1; + block.CodecID = static_cast(ECodec::RAW); + block.OriginalSize = data.size(); + block.OriginalDataRefs = {data}; + + std::vector messages = {{ + .CreatedAt = TInstant::MilliSeconds(123456789), + .MessageMeta = {{"meta-key", "meta-value"}, {std::string(NGrpc::PARTITION_KEY_META_KEY), "partition-key"}}, + .Tx = TTransactionId{.SessionId = "session-id", .TxId = "tx-id"}, + }}; + + const size_t estimate = NGrpc::EstimateTopicWriteRequestBlockSize(block, messages, true); + const size_t actual = BuildTopicWriteRequestBlockSize(block, messages, true); + AssertEstimateCloseToActual(estimate, actual); + } + + Y_UNIT_TEST(EstimateTopicWriteRequestBatchBlockSize) { + TTestBlock block; + block.MessageCount = 3; + block.CodecID = static_cast(ECodec::GZIP); + block.OriginalSize = 100; + block.Data = "compressed-batch"; + block.Compressed = true; + + std::vector messages = { + { + .CreatedAt = TInstant::MilliSeconds(223456789), + .MessageMeta = {{"first-key", "first-value"}}, + .Tx = TTransactionId{.SessionId = "session-id", .TxId = "tx-id"}, + }, + { + .CreatedAt = TInstant::MilliSeconds(323456789), + .MessageMeta = {{std::string(NGrpc::PARTITION_KEY_META_KEY), "partition-key-1"}, {"ignored", "ignored-value"}}, + .Tx = TTransactionId{.SessionId = "session-id", .TxId = "tx-id"}, + }, + { + .CreatedAt = TInstant::MilliSeconds(423456789), + .MessageMeta = {{std::string(NGrpc::PARTITION_KEY_META_KEY), "partition-key-2"}}, + .Tx = TTransactionId{.SessionId = "session-id", .TxId = "tx-id"}, + }, + }; + + const size_t estimate = NGrpc::EstimateTopicWriteRequestBlockSize(block, messages, true); + const size_t actual = BuildTopicWriteRequestBlockSize(block, messages, true); + AssertEstimateCloseToActual(estimate, actual); + } +} + +} // namespace NYdb::NTopic::NTests diff --git a/src/client/types/status/status.cpp b/src/client/types/status/status.cpp index 16ac9da9062..d02530b446c 100644 --- a/src/client/types/status/status.cpp +++ b/src/client/types/status/status.cpp @@ -6,6 +6,8 @@ #include +#include + #include #include diff --git a/src/client/value/value.cpp b/src/client/value/value.cpp index 934aaeb9a32..f525edf1006 100644 --- a/src/client/value/value.cpp +++ b/src/client/value/value.cpp @@ -1437,6 +1437,11 @@ class TValueParser::TImpl { } } + uint32_t GetVariantIndex() const { + CheckKind(ETypeKind::Variant, "GetVariantIndex"); + return GetProto().variant_index(); + } + void CloseVariant() { PopPath(); TypeParser_.CloseVariant(); @@ -2044,6 +2049,10 @@ void TValueParser::OpenVariant() { Impl_->OpenVariant(); } +uint32_t TValueParser::GetVariantIndex() const { + return Impl_->GetVariantIndex(); +} + void TValueParser::CloseVariant() { Impl_->CloseVariant(); } diff --git a/src/library/CMakeLists.txt b/src/library/CMakeLists.txt index 3d3d0a24c60..9df3a0ae027 100644 --- a/src/library/CMakeLists.txt +++ b/src/library/CMakeLists.txt @@ -6,6 +6,7 @@ add_subdirectory(decimal) add_subdirectory(grpc/client) add_subdirectory(issue) add_subdirectory(jwt) +add_subdirectory(kafka) add_subdirectory(malloc/api) add_subdirectory(operation_id) add_subdirectory(persqueue/obfuscate) diff --git a/src/library/kafka/CMakeLists.txt b/src/library/kafka/CMakeLists.txt new file mode 100644 index 00000000000..1cf0913f5c9 --- /dev/null +++ b/src/library/kafka/CMakeLists.txt @@ -0,0 +1,14 @@ +_ydb_sdk_add_library(library-kafka) + +target_link_libraries(library-kafka PUBLIC + yutil + streams-zstd +) + +target_sources(library-kafka PRIVATE + kafka_messages_int.cpp + kafka_records.cpp + kafka_write_buffer.cpp +) + +_ydb_sdk_install_targets(TARGETS library-kafka) diff --git a/src/library/kafka/kafka.h b/src/library/kafka/kafka.h new file mode 100644 index 00000000000..9266872faf3 --- /dev/null +++ b/src/library/kafka/kafka.h @@ -0,0 +1,560 @@ +#pragma once + +#include +#include +#include + +#include + +#include "kafka_write_buffer.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace NKafka { + +enum ECompressionType : ui8 { + NONE = 0, + GZIP = 1, + SNAPPY = 2, + LZ4 = 3, + ZSTD = 4 +}; + +/* + * There are four versions of each field: + * - present version - field serialized and deserialized for this version of protocol. + * - nullable version - field can be null for this versions of protocol. Default field isn't nullable. + * - flexible version - field write as map item tag->value (or tagged version) + * - message flexaible version - version of message support tags + * + * Fields may be of type: + * - bool - fixed length=1 + * - int8 - fixed length=1 + * - int16 - fixed length=2 + * - uint16 - fixed length=2 + * - int32 - fixed length=4 + * - uint32 - fixed length=4 + * - int64 - fixed length=8 + * - uuid - fixed length=16 + * - float64 - fixed length=8 + * - string - can be nullable + * - bytes - can be nullable + * - records - can be nullable + * - struct + * - array - can be nullable + */ + +class TKafkaRecordBatch; + +using TKafkaBool = ui8; +using TKafkaInt8 = i8; +using TKafkaInt16 = i16; +using TKafkaUint16 = ui16; +using TKafkaInt32 = i32; +using TKafkaUint32 = ui32; +using TKafkaInt64 = i64; +using TKafkaUuid = NYdb::NDecimal::TWide; +using TKafkaFloat64 = double; +using TKafkaRawString = TString; +using TKafkaString = std::optional; +using TKafkaRawBytes = TArrayRef; + +inline TKafkaRawBytes ToRawBytes(const TString& str) { + return TKafkaRawBytes(str.data(), str.size()); +} +using TKafkaBytes = std::optional; + +class TKafkaBytesHolder { +public: + TKafkaBytesHolder() = default; + TKafkaBytesHolder(std::nullopt_t) + : Data(std::nullopt) + {} + + TKafkaBytesHolder(TString data) + : Data(std::move(data)) + {} + + TKafkaBytesHolder(TKafkaRawBytes data) + : Data(TString(data.data(), data.size())) + {} + + TKafkaBytesHolder& operator=(std::nullopt_t) { + Data = std::nullopt; + return *this; + } + + TKafkaBytesHolder& operator=(TString data) { + Data = std::move(data); + return *this; + } + + bool has_value() const { + return Data.has_value(); + } + + explicit operator bool() const { + return Data.has_value(); + } + + const TString& value() const { + return Data.value(); + } + + const TString& operator*() const { + return *Data; + } + + const TString* operator->() const { + return &*Data; + } + + bool operator==(const TKafkaBytesHolder& other) const = default; + +private: + std::optional Data; +}; + +using TKafkaRecords = std::optional; + +using TKafkaVersion = i16; + +class TKafkaVersions { +public: + constexpr TKafkaVersions(TKafkaVersion min, TKafkaVersion max) + : Min(min) + , Max(max) { + } + + TKafkaVersion Min; + TKafkaVersion Max; +}; + +static constexpr TKafkaVersions VersionsNever(0, -1); +static constexpr TKafkaVersions VersionsAlways(0, Max()); + +namespace NPrivate { + +struct TKafkaBoolDesc { + static constexpr bool Default = true; + static constexpr bool Nullable = false; + static constexpr bool FixedLength = true; +}; + +struct TKafkaIntDesc { + static constexpr bool Default = true; + static constexpr bool Nullable = false; + static constexpr bool FixedLength = true; +}; + +struct TKafkaVarintDesc { + static constexpr bool Default = true; + static constexpr bool Nullable = false; + static constexpr bool FixedLength = false; +}; + +struct TKafkaUnsignedVarintDesc { + static constexpr bool Default = true; + static constexpr bool Nullable = true; + static constexpr bool FixedLength = false; + + inline static bool IsNull(const TKafkaInt64 value) { return -1 == value; }; +}; + +struct TKafkaUuidDesc { + static constexpr bool Default = true; + static constexpr bool Nullable = false; + static constexpr bool FixedLength = true; +}; + +struct TKafkaFloat64Desc { + static constexpr bool Default = true; + static constexpr bool Nullable = false; + static constexpr bool FixedLength = true; +}; + +struct TKafkaStringDesc { + static constexpr bool Default = true; + static constexpr bool Nullable = true; + static constexpr bool FixedLength = false; + + inline static bool IsNull(const TKafkaString& value) { return !value; }; +}; + +struct TKafkaStructDesc { + static constexpr bool Default = false; + static constexpr bool Nullable = false; + static constexpr bool FixedLength = false; +}; + +struct TKafkaBytesDesc { + static constexpr bool Default = false; + static constexpr bool Nullable = true; + static constexpr bool FixedLength = false; + + inline static bool IsNull(const TKafkaBytes& value) { return !value; }; + inline static bool IsNull(const TKafkaBytesHolder& value) { return !value; }; +}; + +struct TKafkaRecordsDesc { + static constexpr bool Default = false; + static constexpr bool Nullable = false; + static constexpr bool FixedLength = false; +}; + +struct TKafkaArrayDesc { + static constexpr bool Default = false; + static constexpr bool Nullable = true; + static constexpr bool FixedLength = false; + + template + inline static bool IsNull(const std::vector& value) { return value.empty(); }; +}; + +enum ESizeFormat { + Default = 0, + Varint = 1 +}; + +} // namespace NPrivate + + +// see https://kafka.apache.org/11/protocol.html#protocol_error_codes +enum EKafkaErrors { + + UNKNOWN_SERVER_ERROR = -1, // The server experienced an unexpected error when processing the request. + NONE_ERROR = 0, + OFFSET_OUT_OF_RANGE = 1, // The requested offset is not within the range of offsets maintained by the server., + CORRUPT_MESSAGE = 2, // This message has failed its CRC checksum, exceeds the valid size, has a null key for a compacted topic, or is otherwise corrupt. + UNKNOWN_TOPIC_OR_PARTITION = 3, // This server does not host this topic-partition. + INVALID_FETCH_SIZE = 4, // The requested fetch size is invalid. + LEADER_NOT_AVAILABLE = 5, // There is no leader for this topic-partition as we are in the middle of a leadership election. + NOT_LEADER_OR_FOLLOWER = 6, // For requests intended only for the leader, this error indicates that the broker is not the current leader. + // For requests intended for any replica, this error indicates that the broker is not a replica of the topic partition. + REQUEST_TIMED_OUT = 7, // The request timed out. + BROKER_NOT_AVAILABLE = 8, // The broker is not available. + REPLICA_NOT_AVAILABLE = 9, // The replica is not available for the requested topic-partition. Produce/Fetch requests and other requests + // intended only for the leader or follower return NOT_LEADER_OR_FOLLOWER if the broker is not a replica of the topic-partition. + MESSAGE_TOO_LARGE = 10, // The request included a message larger than the max message size the server will accept. + STALE_CONTROLLER_EPOCH = 11, // The controller moved to another broker. + OFFSET_METADATA_TOO_LARGE = 12, // The metadata field of the offset request was too large. + NETWORK_EXCEPTION = 13, // The server disconnected before a response was received. + COORDINATOR_LOAD_IN_PROGRESS = 14, // The coordinator is loading and hence can't process requests. + COORDINATOR_NOT_AVAILABLE = 15, // The coordinator is not available. + NOT_COORDINATOR = 16, // This is not the correct coordinator. + INVALID_TOPIC_EXCEPTION = 17, // The request attempted to perform an operation on an invalid topic. + RECORD_LIST_TOO_LARGE = 18, // The request included message batch larger than the configured segment size on the server. + NOT_ENOUGH_REPLICAS = 19, // Messages are rejected since there are fewer in-sync replicas than required. + NOT_ENOUGH_REPLICAS_AFTER_APPEND = 20, // Messages are written to the log, but to fewer in-sync replicas than required. + INVALID_REQUIRED_ACKS = 21, // Produce request specified an invalid value for required acks. + ILLEGAL_GENERATION = 22, // Specified group generation id is not valid. + INCONSISTENT_GROUP_PROTOCOL = 23, // The group member's supported protocols are incompatible with those of existing members + // or first group member tried to join with empty protocol type or empty protocol list. + INVALID_GROUP_ID = 24, // The configured groupId is invalid. + UNKNOWN_MEMBER_ID = 25, // The coordinator is not aware of this member. + INVALID_SESSION_TIMEOUT = 26, // The session timeout is not within the range allowed by the broker + // (as configured by group.min.session.timeout.ms and group.max.session.timeout.ms). + REBALANCE_IN_PROGRESS = 27, // The group is rebalancing, so a rejoin is needed. + INVALID_COMMIT_OFFSET_SIZE = 28, // The committing offset data size is not valid. + TOPIC_AUTHORIZATION_FAILED = 29, // Topic authorization failed. + GROUP_AUTHORIZATION_FAILED = 30, // Group authorization failed. + CLUSTER_AUTHORIZATION_FAILED = 31, // Cluster authorization failed. + INVALID_TIMESTAMP = 32, // The timestamp of the message is out of acceptable range. + UNSUPPORTED_SASL_MECHANISM = 33, // The broker does not support the requested SASL mechanism. + ILLEGAL_SASL_STATE = 34, // Request is not valid given the current SASL state. + UNSUPPORTED_VERSION = 35, // The version of API is not supported. + TOPIC_ALREADY_EXISTS = 36, // Topic with this name already exists. + INVALID_PARTITIONS = 37, // Number of partitions is below 1. + INVALID_REPLICATION_FACTOR = 38, // Replication factor is below 1 or larger than the number of available brokers. + INVALID_REPLICA_ASSIGNMENT = 39, // Replica assignment is invalid. + INVALID_CONFIG = 40, // Configuration is invalid. + NOT_CONTROLLER = 41, // This is not the correct controller for this cluster. + INVALID_REQUEST = 42, // This most likely occurs because of a request being malformed by the + // client library or the message was sent to an incompatible broker. See the broker logs + // for more details. + UNSUPPORTED_FOR_MESSAGE_FORMAT = 43, // The message format version on the broker does not support the request. + POLICY_VIOLATION = 44, // Request parameters do not satisfy the configured policy. + OUT_OF_ORDER_SEQUENCE_NUMBER = 45, // The broker received an out of order sequence number. + DUPLICATE_SEQUENCE_NUMBER = 46, // The broker received a duplicate sequence number. + INVALID_PRODUCER_EPOCH = 47, // Producer attempted to produce with an old epoch. + INVALID_TXN_STATE = 48, // The producer attempted a transactional operation in an invalid state. + INVALID_PRODUCER_ID_MAPPING = 49, // The producer attempted to use a producer id which is not currently assigned to + // its transactional id. + INVALID_TRANSACTION_TIMEOUT = 50, // The transaction timeout is larger than the maximum value allowed by + // the broker (as configured by transaction.max.timeout.ms). + CONCURRENT_TRANSACTIONS = 51, // The producer attempted to update a transaction + // while another concurrent operation on the same transaction was ongoing. + TRANSACTION_COORDINATOR_FENCED = 52, // Indicates that the transaction coordinator sending a WriteTxnMarker + // is no longer the current coordinator for a given producer. + TRANSACTIONAL_ID_AUTHORIZATION_FAILED = 53, // Transactional Id authorization failed. + SECURITY_DISABLED = 54, // Security features are disabled. + OPERATION_NOT_ATTEMPTED = 55, // The broker did not attempt to execute this operation. This may happen for + // batched RPCs where some operations in the batch failed, causing the broker to respond without + // trying the rest. + KAFKA_STORAGE_ERROR = 56, // Disk error when trying to access log file on the disk. + LOG_DIR_NOT_FOUND = 57, // The user-specified log directory is not found in the broker config. + SASL_AUTHENTICATION_FAILED = 58, // SASL Authentication failed. + UNKNOWN_PRODUCER_ID = 59, // This exception is raised by the broker if it could not locate the producer metadata + // associated with the producerId in question. This could happen if, for instance, the producer's records + // were deleted because their retention time had elapsed. Once the last records of the producerId are + // removed, the producer's metadata is removed from the broker, and future appends by the producer will + // return this exception. + REASSIGNMENT_IN_PROGRESS = 60, // A partition reassignment is in progress. + DELEGATION_TOKEN_AUTH_DISABLED = 61, // Delegation Token feature is not enabled. + DELEGATION_TOKEN_NOT_FOUND = 62, // Delegation Token is not found on server. + DELEGATION_TOKEN_OWNER_MISMATCH = 63, // Specified Principal is not valid Owner/Renewer. + DELEGATION_TOKEN_REQUEST_NOT_ALLOWED = 64, // Delegation Token requests are not allowed on PLAINTEXT/1-way SSL + // channels and on delegation token authenticated channels. + DELEGATION_TOKEN_AUTHORIZATION_FAILED = 65, // Delegation Token authorization failed. + DELEGATION_TOKEN_EXPIRED = 66, // Delegation Token is expired. + INVALID_PRINCIPAL_TYPE = 67, // Supplied principalType is not supported. + NON_EMPTY_GROUP = 68, // The group is not empty. + GROUP_ID_NOT_FOUND = 69, // The group id does not exist. + FETCH_SESSION_ID_NOT_FOUND = 70, // The fetch session ID was not found. + INVALID_FETCH_SESSION_EPOCH = 71, // The fetch session epoch is invalid. + LISTENER_NOT_FOUND = 72, // There is no listener on the leader broker that matches the listener on which + // metadata request was processed. + TOPIC_DELETION_DISABLED = 73, // Topic deletion is disabled. + FENCED_LEADER_EPOCH = 74, // The leader epoch in the request is older than the epoch on the broker. + UNKNOWN_LEADER_EPOCH = 75, // The leader epoch in the request is newer than the epoch on the broker. + UNSUPPORTED_COMPRESSION_TYPE = 76, // The requesting client does not support the compression type of given partition. + STALE_BROKER_EPOCH = 77, // Broker epoch has changed. + OFFSET_NOT_AVAILABLE = 78, // The leader high watermark has not caught up from a recent leader + // election so the offsets cannot be guaranteed to be monotonically increasing. + MEMBER_ID_REQUIRED = 79, // The group member needs to have a valid member id before actually entering a consumer group. + PREFERRED_LEADER_NOT_AVAILABLE = 80, // The preferred leader was not available. + GROUP_MAX_SIZE_REACHED = 81, // The consumer group has reached its max size. + FENCED_INSTANCE_ID = 82, // The broker rejected this static consumer since + // another consumer with the same group.instance.id has registered with a different member.id. + ELIGIBLE_LEADERS_NOT_AVAILABLE = 83, // Eligible topic partition leaders are not available. + ELECTION_NOT_NEEDED = 84, // Leader election not needed for topic partition. + NO_REASSIGNMENT_IN_PROGRESS = 85, // No partition reassignment is in progress. + GROUP_SUBSCRIBED_TO_TOPIC = 86, // Deleting offsets of a topic is forbidden while the consumer group is actively subscribed to it. + INVALID_RECORD = 87, // This record has failed the validation on broker and hence will be rejected. + UNSTABLE_OFFSET_COMMIT = 88, // There are unstable offsets that need to be cleared. + THROTTLING_QUOTA_EXCEEDED = 89, // The throttling quota has been exceeded. + PRODUCER_FENCED = 90, // There is a newer producer with the same transactionalId + // which fences the current one. + RESOURCE_NOT_FOUND = 91, // A request illegally referred to a resource that does not exist. + DUPLICATE_RESOURCE = 92, // A request illegally referred to the same resource twice. + UNACCEPTABLE_CREDENTIAL = 93, // Requested credential would not meet criteria for acceptability. + INCONSISTENT_VOTER_SET = 94, // Indicates that the either the sender or recipient of a + // voter-only request is not one of the expected voters + INVALID_UPDATE_VERSION = 95, // The given update version was invalid. + FEATURE_UPDATE_FAILED = 96, // Unable to update finalized features due to an unexpected server error. + PRINCIPAL_DESERIALIZATION_FAILURE = 97, // Request principal deserialization failed during forwarding. + // This indicates an internal error on the broker cluster security setup. + SNAPSHOT_NOT_FOUND = 98, // Requested snapshot was not found + POSITION_OUT_OF_RANGE = 99, // Requested position is not greater than or equal to zero, and less than the size of the snapshot. + UNKNOWN_TOPIC_ID = 100, // This server does not host this topic ID. + DUPLICATE_BROKER_REGISTRATION = 101, // This broker ID is already in use. + BROKER_ID_NOT_REGISTERED = 102, // The given broker ID was not registered. + INCONSISTENT_TOPIC_ID = 103, // The log's topic ID did not match the topic ID in the request + INCONSISTENT_CLUSTER_ID = 104, // The clusterId in the request does not match that found on the server + TRANSACTIONAL_ID_NOT_FOUND = 105, // The transactionalId could not be found + FETCH_SESSION_TOPIC_ID_ERROR = 106, // The fetch session encountered inconsistent topic ID usage + INELIGIBLE_REPLICA = 107, // The new ISR contains at least one ineligible replica. + NEW_LEADER_ELECTED = 108 // The AlterPartition request successfully updated the partition state but the leader has changed. +}; + +template +void NormalizeNumber(T& value) { +#ifndef WORDS_BIGENDIAN + char* b = (char*)&value; + char* e = b + sizeof(T) - 1; + while (b < e) { + std::swap(*b, *e); + ++b; + --e; + } +#endif +} + +template, typename U = std::make_unsigned_t> +U AsUnsigned(S value) { + static constexpr ui8 Shift = (sizeof(T) << 3) - 1; + return (value << 1) ^ (value >> Shift); +} + +class TKafkaWritable { +public: + TKafkaWritable(TKafkaWriteBuffer& buffer) + : Buffer(buffer){}; + + template + TKafkaWritable& operator<<(const T val) { + NormalizeNumber(val); + write((const char*)&val, sizeof(T)); + return *this; + }; + + TKafkaWritable& operator<<(const TKafkaUuid& val); + TKafkaWritable& operator<<(const TKafkaRawBytes& val); + TKafkaWritable& operator<<(const TKafkaRawString& val); + + template> + void writeUnsignedVarint(T v) { + static constexpr U Mask = Max() - 0x7F; + + U value = v; + while ((value & Mask) != 0L) { + ui8 b = (ui8) ((value & 0x7f) | 0x80); + write((const char*)&b, sizeof(b)); + value >>= 7; + } + ui8 b = (ui8) value; + write((const char*)&b, sizeof(b)); + } + + template> + void writeVarint(T value) { + writeUnsignedVarint(AsUnsigned(value)); + } + + void write(const char* val, size_t length); + +private: + TKafkaWriteBuffer& Buffer; +}; + +class TKafkaReadable { +public: + TKafkaReadable(const TBuffer& is) + : Is(is) + , Position(0) { + } + + template + TKafkaReadable& operator>>(T& val) { + char* v = (char*)&val; + read(v, sizeof(T)); + NormalizeNumber(val); + return *this; + }; + + TKafkaReadable& operator>>(TKafkaUuid& val); + + void read(char* val, size_t length); + char get(); + + template + U readUnsignedVarint() { + static constexpr size_t MaxLength = (sizeof(U) * 8 - 1) / 7 * 7; + + U value = 0; + size_t i = 0; + U b; + while (((b = static_cast(get())) & 0x80) != 0) { + if (i >= MaxLength) { + ythrow yexception() << "illegal varint length"; + } + value |= (b & 0x7f) << i; + i += 7; + } + + value |= b << i; + return value; + } + + template> + S readVarint() { + U v = readUnsignedVarint(); + return (v >> 1) ^ -static_cast(v & 1); + } + + TArrayRef Bytes(size_t length); + + // returns a character from the specified position. The current position does not change. + char take(size_t shift); + + void skip(size_t length); + + size_t left() const; + + size_t position() const; + + void SetAllowCompressed(bool allowCompressed) { + AllowCompressed_ = allowCompressed; + } + + bool GetAllowCompressed() const { + return AllowCompressed_; + } + +private: + void checkEof(size_t length); + + const TBuffer& Is; + size_t Position; + // Temporary switch until server-side Kafka record batch support is implemented. + bool AllowCompressed_ = false; +}; + +struct TReadDemand { + constexpr TReadDemand() + : Buffer(nullptr) + , Length(0) { + } + + constexpr TReadDemand(char* buffer, size_t length) + : Buffer(buffer) + , Length(length) { + } + + constexpr TReadDemand(size_t length) + : Buffer(nullptr) + , Length(length) { + } + + char* GetBuffer() const { + return Buffer; + } + size_t GetLength() const { + return Length; + } + explicit operator bool() const { + return 0 < Length; + } + bool Skip() const { + return nullptr == Buffer; + } + + char* Buffer; + size_t Length; +}; + +static constexpr TReadDemand NoDemand; + +class TMessage { +public: + virtual ~TMessage() = default; + + virtual i32 Size(TKafkaVersion version) const = 0; + virtual void Read(TKafkaReadable& readable, TKafkaVersion version) = 0; + virtual void Write(TKafkaWritable& writable, TKafkaVersion version) const = 0; + + bool operator==(const TMessage& other) const = default; +}; + +class TApiMessage: public TMessage { +public: + using TPtr = std::shared_ptr; + + ~TApiMessage() = default; + + virtual i16 ApiKey() const = 0; +}; + +std::unique_ptr CreateRequest(i16 apiKey); +std::unique_ptr CreateResponse(i16 apiKey); + +i16 RequestHeaderVersion(i16 apiKey, TKafkaVersion version); +i16 ResponseHeaderVersion(i16 apiKey, TKafkaVersion version); + +} // namespace NKafka diff --git a/src/library/kafka/kafka_log.h b/src/library/kafka/kafka_log.h new file mode 100644 index 00000000000..6e67e53982c --- /dev/null +++ b/src/library/kafka/kafka_log.h @@ -0,0 +1,11 @@ +#pragma once + +#include + +namespace NKafka { + +static constexpr bool DEBUG_ENABLED = false; + +TString Hex(const char* begin, const char* end); + +} diff --git a/src/library/kafka/kafka_messages_int.cpp b/src/library/kafka/kafka_messages_int.cpp new file mode 100644 index 00000000000..ecfcc6cf955 --- /dev/null +++ b/src/library/kafka/kafka_messages_int.cpp @@ -0,0 +1,98 @@ +#include "kafka_messages_int.h" + +#include + +namespace NKafka { + +TKafkaWritable& TKafkaWritable::operator<<(const TKafkaRawBytes& val) { + write(val.data(), val.size()); + return *this; +} + +TKafkaWritable& TKafkaWritable::operator<<(const TKafkaRawString& val) { + write(val.data(), val.length()); + return *this; +} + +TKafkaWritable& TKafkaWritable::operator<<(const TKafkaUuid& val) { + ui64 h = ui64(val >> (sizeof(ui64) << 3)); + ui64 l = ui64(val); + *this << h << l; + return *this; +} + +void TKafkaWritable::write(const char* val, size_t length) { + Buffer.write(val, length); +} + +TKafkaReadable& TKafkaReadable::operator>>(TKafkaUuid& val) { + ui64 h; + ui64 l; + + *this >> h >> l; + + val = TKafkaUuid(h, l); + return *this; +} + + +void TKafkaReadable::read(char* val, size_t length) { + checkEof(length); + memcpy(val, Is.Data() + Position, length); + Position += length; +} + +char TKafkaReadable::get() { + char r; + read(&r, sizeof(r)); + return r; +} + +TArrayRef TKafkaReadable::Bytes(size_t length) { + checkEof(length); + TArrayRef r(Is.Data() + Position, length); + Position += length; + return r; +} + +void TKafkaReadable::skip(size_t length) { + checkEof(length); + Position += length; +} + +char TKafkaReadable::take(size_t shift) { + checkEof(shift + sizeof(char)); + return *(Is.Data() + Position + shift); +} + +size_t TKafkaReadable::left() const { + return Is.Size() - Position; +} + +size_t TKafkaReadable::position() const { + return Position; +} + +void TKafkaReadable::checkEof(size_t length) { + if (Position + length > Is.Size()) { + ythrow yexception() << "unexpected end of stream"; + } +} + +char Hex(const unsigned char c) { + return c < 10 ? '0' + c : 'A' + c - 10; +} + +TString Hex(const char* begin, const char *end) { + TStringBuilder sb; + for(auto i = begin; i < end; ++i) { + unsigned char c = *i; + if (i != begin) { + sb << ", "; + } + sb << "0x" << Hex(c >> 4) << Hex(c & 0x0F); + } + return sb; +} + +} // namespace NKafka diff --git a/src/library/kafka/kafka_messages_int.h b/src/library/kafka/kafka_messages_int.h new file mode 100644 index 00000000000..99469c65fa1 --- /dev/null +++ b/src/library/kafka/kafka_messages_int.h @@ -0,0 +1,730 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +#include + +#include "kafka_records.h" +#include "kafka_log.h" + +namespace NKafka { +namespace NPrivate { + +struct TWriteCollector { + ui32 NumTaggedFields = 0; +}; + +struct TSizeCollector { + ui32 Size = 0; + ui32 NumTaggedFields = 0; +}; + +template> +size_t SizeOfUnsignedVarint(T v) { + static constexpr U Mask = Max() - 0x7F; + + U value = v; + size_t bytes = 1; + while ((value & Mask) != 0L) { + bytes += 1; + value >>= 7; + } + return bytes; +} + +template +size_t SizeOfVarint(T value) { + return SizeOfUnsignedVarint(AsUnsigned(value)); +} + +template +constexpr bool VersionAll() { + return 0 == min && max == Max(); +} + +template +constexpr bool VersionNone() { + return 0 == min && max == -1; +} + +template +inline bool VersionCheck(const TKafkaVersion value) { + if constexpr (VersionNone()) { + return false; + } else if constexpr (VersionAll()) { + return true; + } else if constexpr (max == Max()) { + return min <= value; + } else if constexpr (min == 0) { + return value <= max; + } else { + return min <= value && value <= max; + } +} + +template +bool IsDefaultValue(const typename Meta::Type& value) { + if constexpr (std::is_base_of_v) { + typename Meta::Type defValue; + return defValue == value; + } else if constexpr (Meta::TypeDesc::Default) { + return Meta::Default == value; + } else if constexpr (Meta::TypeDesc::Nullable) { + return Meta::TypeDesc::IsNull(value); + } else { + return false; + } +} + + +template +struct HasSizeFormat: std::false_type {}; + +template +struct HasSizeFormat: std::true_type {}; + +template +constexpr ESizeFormat SizeFormat() { + if constexpr (HasSizeFormat::value) { + return Meta::SizeFormat; + } else { + return ESizeFormat::Default; + } +} + + +template +inline void WriteStringSize(TKafkaWritable& writable, TKafkaVersion version, TKafkaInt32 value) { + if (VersionCheck(version)) { + writable.writeUnsignedVarint(value + 1); + } else { + writable << (TKafkaInt16)value; + } +} + +template +inline TKafkaInt32 ReadStringSize(TKafkaReadable& readable, TKafkaVersion version) { + if (VersionCheck(version)) { + return readable.readUnsignedVarint() - 1; + } else { + TKafkaInt16 v; + readable >> v; + return v; + } +} + +template +inline void WriteArraySize(TKafkaWritable& writable, TKafkaVersion version, TKafkaInt32 value) { + if constexpr (SizeFormat() == Varint) { + return writable.writeVarint(value); + } else if (VersionCheck(version)) { + writable.writeUnsignedVarint(value + 1); + } else { + writable << value; + } +} + +template +inline TKafkaInt32 ReadArraySize(TKafkaReadable& readable, TKafkaVersion version) { + if constexpr (SizeFormat() == Varint) { + return readable.readVarint(); + } else if (VersionCheck(version)) { + return readable.readUnsignedVarint() - 1; + } else { + TKafkaInt32 v; + readable >> v; + return v; + } +} + +template +inline TKafkaInt32 ArraySize(TKafkaVersion version, TKafkaInt32 size) { + if constexpr (SizeFormat() == Varint) { + return SizeOfVarint(size); + } else if (VersionCheck(version)) { + return SizeOfUnsignedVarint(size + 1); + } else { + return sizeof(TKafkaInt32); + } +} + + +inline IOutputStream& operator <<(IOutputStream& out, const TKafkaUuid& /*value*/) { + return out << "---"; +} + + + +// +// Common +// +template +class TypeStrategy { +public: + inline static void DoWrite(TKafkaWritable& writable, TKafkaVersion /*version*/, const TValueType& value) { + writable << value; + } + + inline static void DoWriteTag(TKafkaWritable& writable, TKafkaVersion /*version*/, const TValueType& value) { + writable << value; + } + + inline static void DoRead(TKafkaReadable& readable, TKafkaVersion /*version*/, TValueType& value) { + readable >> value; + } + + inline static i64 DoSize(TKafkaVersion /*version*/, const TValueType& /*value*/) { + static_assert(TTypeDesc::FixedLength, "Unsupported type: serialization length must be constant"); + return sizeof(TValueType); + } + + inline static void DoLog(const TValueType& value) { + if constexpr (DEBUG_ENABLED) { + if constexpr (sizeof(TValueType) == 1) { + Cerr << "Was read field '" << Meta::Name << "' value " << (size_t)value << Endl; + } else { + Cerr << "Was read field '" << Meta::Name << "' value " << value << Endl; + } + } + } +}; + + +// +// TKafkaVarintDesc +// +template +class TypeStrategy { +public: + inline static void DoWrite(TKafkaWritable& writable, TKafkaVersion version, const TValueType& value) { + if (VersionCheck(version)) { + writable.writeVarint(value); + } else { + writable << value; + } + } + + inline static void DoWriteTag(TKafkaWritable& writable, TKafkaVersion version, const TValueType& value) { + if (VersionCheck(version)) { + writable.writeVarint(value); + } else { + writable << value; + } + } + + inline static void DoRead(TKafkaReadable& readable, TKafkaVersion version, TValueType& value) { + if (VersionCheck(version)) { + value = readable.readVarint(); + } else { + readable >> value; + } + } + + inline static i64 DoSize(TKafkaVersion version, const TValueType& value) { + if (VersionCheck(version)) { + return SizeOfVarint(value); + } else { + return sizeof(TValueType); + } + } + + inline static void DoLog(const TValueType& value) { + if constexpr (DEBUG_ENABLED) { + Cerr << "Was read field '" << Meta::Name << "' value " << value << Endl; + } + } +}; + + +// +// TKafkaUnsignedVarintDesc +// +template +class TypeStrategy { +public: + inline static void DoWrite(TKafkaWritable& writable, TKafkaVersion version, const TValueType& value) { + if (VersionCheck(version)) { + writable.writeUnsignedVarint(value + 1); + } else { + writable << value; + } + } + + inline static void DoWriteTag(TKafkaWritable& writable, TKafkaVersion version, const TValueType& value) { + if (VersionCheck(version)) { + writable.writeUnsignedVarint(value + 1); + } else { + writable << value; + } + } + + inline static void DoRead(TKafkaReadable& readable, TKafkaVersion version, TValueType& value) { + if (VersionCheck(version)) { + value = readable.readUnsignedVarint() - 1; + } else { + readable >> value; + } + } + + inline static i64 DoSize(TKafkaVersion version, const TValueType& value) { + if (VersionCheck(version)) { + return SizeOfUnsignedVarint(value); + } else { + return sizeof(TValueType); + } + } + + inline static void DoLog(const TValueType& value) { + if constexpr (DEBUG_ENABLED) { + Cerr << "Was read field '" << Meta::Name << "' value " << value << Endl; + } + } +}; + + +// +// TKafkaStructDesc +// +template +class TypeStrategy { +public: + inline static void DoWrite(TKafkaWritable& writable, TKafkaVersion version, const TValueType& value) { + value.Write(writable, version); + } + + inline static void DoWriteTag(TKafkaWritable& writable, TKafkaVersion version, const TValueType& value) { + value.Write(writable, version); + } + + inline static void DoRead(TKafkaReadable& readable, TKafkaVersion version, TValueType& value) { + value.Read(readable, version); + } + + inline static i64 DoSize(TKafkaVersion version, const TValueType& value) { + return value.Size(version); + } + + inline static void DoLog(const TValueType& /*value*/) { + if constexpr (DEBUG_ENABLED) { + Cerr << "Was read field '" << Meta::Name << "' type struct" << Endl; + } + } +}; + + +// +// TKafkaString +// +template +class TypeStrategy { +public: + inline static void DoWrite(TKafkaWritable& writable, TKafkaVersion version, const TKafkaString& value) { + if (value) { + const auto& v = *value; + if (v.length() > Max()) { + ythrow yexception() << "string field " << Meta::Name << " is too long to be serialized " << v.length(); + } + WriteStringSize(writable, version, v.length()); + writable << v; + } else { + if (VersionCheck(version)) { + WriteStringSize(writable, version, -1); + } else { + ythrow yexception() << "non-nullable field " << Meta::Name << " serializing as null"; + } + } + } + + inline static void DoWriteTag(TKafkaWritable& writable, TKafkaVersion version, const TKafkaString& value) { + const auto& v = *value; + WriteStringSize(writable, version, v.length()); + writable << v; + } + + inline static void DoRead(TKafkaReadable& readable, TKafkaVersion version, TKafkaString& value) { + TKafkaInt32 length = ReadStringSize(readable, version); + if (length < 0) { + if (VersionCheck(version)) { + value = std::nullopt; + } else { + ythrow yexception() << "non-nullable field " << Meta::Name << " was serialized as null"; + } + } else if (length > Max()){ + ythrow yexception() << "string field " << Meta::Name << " had invalid length " << length; + } else { + value = TString(); + value->ReserveAndResize(length); + readable.read(const_cast(value->data()), length); + } + } + + inline static i64 DoSize(TKafkaVersion version, const TKafkaString& value) { + if (value) { + const auto& v = *value; + if (v.length() > Max()) { + ythrow yexception() << "string field " << Meta::Name << " is too long to be serialized " << v.length(); + } + if (VersionCheck(version)) { + return v.length() + SizeOfUnsignedVarint(v.length() + 1); + } else { + return v.length() + sizeof(TKafkaInt16); + } + } else { + if (VersionCheck(version)) { + return sizeof(TKafkaInt8); + } else { + return sizeof(TKafkaInt16); + } + } + } + + inline static void DoLog(const TKafkaString& value) { + if constexpr (DEBUG_ENABLED) { + Cerr << "Was read field '" << Meta::Name << "' type String. Size " << (value ? value->size() : 0) << " Value: " << (value ? *value : "null") << Endl; + } + } +}; + + +// +// TKafkaBytes +// +template +class TypeStrategy { +public: + inline static void DoWrite(TKafkaWritable& writable, TKafkaVersion version, const TKafkaBytes& value) { + if (value) { + const auto& v = *value; + WriteArraySize(writable, version, v.size()); + writable.write(v.data(), v.size()); + } else { + if (VersionCheck(version)) { + WriteArraySize(writable, version, -1); + } else { + ythrow yexception() << "non-nullable field " << Meta::Name << " serializing as null"; + } + } + } + + inline static void DoWriteTag(TKafkaWritable& writable, TKafkaVersion version, const TKafkaBytes& value) { + const auto& v = *value; + WriteArraySize(writable, version, v.size()); + writable.write(v.data(), v.size()); + } + + inline static void DoRead(TKafkaReadable& readable, TKafkaVersion version, TKafkaBytes& value) { + TKafkaInt32 length = ReadArraySize(readable, version); + if (length < 0) { + if (VersionCheck(version)) { + value = std::nullopt; + } else { + ythrow yexception() << "non-nullable field " << Meta::Name << " was serialized as null"; + } + } else { + value = readable.Bytes(length); + } + } + + inline static i64 DoSize(TKafkaVersion version, const TKafkaBytes& value) { + if (value) { + const auto& v = *value; + return v.size() + ArraySize(version, v.size()); + } else { + if (VersionCheck(version)) { + return 1; + } else { + return sizeof(TKafkaInt32); + } + } + } + + inline static void DoLog(const TKafkaBytes& value) { + if constexpr (DEBUG_ENABLED) { + Cerr << "Was read field '" << Meta::Name << "' type Bytes. Size " << (value ? value->size() : 0) << Endl; + } + } +}; + +template +class TypeStrategy { +public: + inline static void DoWrite(TKafkaWritable& writable, TKafkaVersion version, const TKafkaBytesHolder& value) { + if (value) { + const auto& v = *value; + WriteArraySize(writable, version, v.size()); + writable.write(v.data(), v.size()); + } else { + if (VersionCheck(version)) { + WriteArraySize(writable, version, -1); + } else { + ythrow yexception() << "non-nullable field " << Meta::Name << " serializing as null"; + } + } + } + + inline static void DoWriteTag(TKafkaWritable& writable, TKafkaVersion version, const TKafkaBytesHolder& value) { + const auto& v = *value; + WriteArraySize(writable, version, v.size()); + writable.write(v.data(), v.size()); + } + + inline static void DoRead(TKafkaReadable& readable, TKafkaVersion version, TKafkaBytesHolder& value) { + TKafkaInt32 length = ReadArraySize(readable, version); + if (length < 0) { + if (VersionCheck(version)) { + value = std::nullopt; + } else { + ythrow yexception() << "non-nullable field " << Meta::Name << " was serialized as null"; + } + } else { + TString data; + data.ReserveAndResize(length); + readable.read(const_cast(data.data()), length); + value = std::move(data); + } + } + + inline static i64 DoSize(TKafkaVersion version, const TKafkaBytesHolder& value) { + if (value) { + return value->size() + ArraySize(version, value->size()); + } else { + if (VersionCheck(version)) { + return 1; + } else { + return sizeof(TKafkaInt32); + } + } + } + + inline static void DoLog(const TKafkaBytesHolder& value) { + if constexpr (DEBUG_ENABLED) { + Cerr << "Was read field '" << Meta::Name << "' type BytesHolder. Size " << (value ? value->size() : 0) << Endl; + } + } +}; + + +// +// TKafkaRecords +// +template +class TypeStrategy { + static constexpr TKafkaVersion CURRENT_RECORD_VERSION = 2; +public: + inline static void DoWrite(TKafkaWritable& writable, TKafkaVersion version, const TKafkaRecords& value) { + if (value) { + WriteArraySize(writable, version, value->Size(CURRENT_RECORD_VERSION)); + (*value).Write(writable, CURRENT_RECORD_VERSION); + } else { + WriteArraySize(writable, version, 0); + } + } + + inline static void DoWriteTag(TKafkaWritable& writable, TKafkaVersion version, const TKafkaRecords& value) { + DoWrite(writable, version, value); + } + + inline static void DoRead(TKafkaReadable& readable, TKafkaVersion version, TKafkaRecords& value) { + int length = ReadArraySize(readable, version); + if (length > 0) { + char magic = readable.take(16); + value.emplace(); + + if (magic < CURRENT_RECORD_VERSION) { + ReadLegacyRecordBatch(readable, magic, length, *value); + } else { + (*value).Read(readable, magic); + } + } else { + value = std::nullopt; + } + } + + inline static i64 DoSize(TKafkaVersion version, const TKafkaRecords& value) { + if (value) { + const auto& v = *value; + const auto size = v.Size(CURRENT_RECORD_VERSION); + return size + ArraySize(version, size); + } else { + if (VersionCheck(version)) { + return 1; + } else { + return sizeof(TKafkaInt32); + } + } + } + + inline static void DoLog(const TKafkaRecords& /*value*/) { + if constexpr (DEBUG_ENABLED) { + Cerr << "Was read field '" << Meta::Name << "' type Records." << Endl; + } + } +}; + + +// +// KafkaArray +// +template +class TypeStrategy, TKafkaArrayDesc> { + using ItemStrategy = TypeStrategy; + +public: + inline static void DoWrite(TKafkaWritable& writable, TKafkaVersion version, const std::vector& value) { + WriteArraySize(writable, version, value.size()); + + for(const auto& v : value) { + ItemStrategy::DoWrite(writable, version, v); + } + } + + inline static void DoWriteTag(TKafkaWritable& writable, TKafkaVersion version, const std::vector& value) { + WriteArraySize(writable, version, value.size()); + + for(const auto& v : value) { + ItemStrategy::DoWrite(writable, version, v); + } + } + + inline static void DoRead(TKafkaReadable& readable, TKafkaVersion version, std::vector& value) { + TKafkaInt32 length = ReadArraySize(readable, version); + if (length < 0) { + if (VersionCheck(version)) { + return; + } else { + ythrow yexception() << "non-nullable field " << Meta::Name << " was serialized as null"; + } + } + value.resize(length); + + for (int i = 0; i < length; ++i) { + ItemStrategy::DoRead(readable, version, value[i]); + } + } + + inline static i64 DoSize(TKafkaVersion version, const std::vector& value) { + TKafkaInt32 size = 0; + if constexpr (Meta::TypeDesc::FixedLength) { + size = value.size() * sizeof(TValueType); + } else { + for(const auto& v : value) { + size += ItemStrategy::DoSize(version, v); + } + } + return size + ArraySize(version, value.size()); + } + + inline static void DoLog(const std::vector& value) { + if constexpr (DEBUG_ENABLED) { + Cerr << "Was read field '" << Meta::Name << "' type Array. Size " << value.size() << Endl; + } + } +}; + + + +// +// Main fields function +// +template +inline void Write(TWriteCollector& collector, TKafkaWritable& writable, TKafkaInt16 version, const typename Meta::Type& value) { + if (VersionCheck(version)) { + if (VersionCheck(version)) { + if (!IsDefaultValue(value)) { + ++collector.NumTaggedFields; + } + } else { + TypeStrategy::DoWrite(writable, version, value); + } + } +} + +template +inline void Read(TKafkaReadable& readable, TKafkaInt16 version, typename Meta::Type& value) { + if (!VersionNone() + && VersionCheck(version)) { + return; + } else { + if (VersionCheck(version)) { + try { + TypeStrategy::DoRead(readable, version, value); + TypeStrategy::DoLog(value); + } catch (const yexception& e) { + ythrow yexception() << "error on read field " << Meta::Name << ": " << e.what(); + } + } else if constexpr (Meta::TypeDesc::Default) { + value = Meta::Default; + } + } +} + +template +inline void Size(TSizeCollector& collector, TKafkaInt16 version, const typename Meta::Type& value) { + if constexpr (!VersionNone()) { + if (VersionCheck(version)) { + if (VersionCheck(version)) { + if (!IsDefaultValue(value)) { + ++collector.NumTaggedFields; + + i64 size = TypeStrategy::DoSize(version, value); + collector.Size += size + SizeOfUnsignedVarint(Meta::Tag) + SizeOfUnsignedVarint(size); + if constexpr (DEBUG_ENABLED) { + Cerr << "Size of field '" << Meta::Name << "' " << size << " + " << SizeOfUnsignedVarint(Meta::Tag) << " + " << SizeOfUnsignedVarint(size) << Endl; + } + } + } else { + i64 size = TypeStrategy::DoSize(version, value); + collector.Size += size; + if constexpr (DEBUG_ENABLED) { + Cerr << "Size of field '" << Meta::Name << "' " << size << Endl; + } + } + } + } else { + if (VersionCheck(version)) { + i64 size = TypeStrategy::DoSize(version, value); + collector.Size += size; + if constexpr (DEBUG_ENABLED) { + Cerr << "Size of field '" << Meta::Name << "' " << size << Endl; + } + } + } +} + +template +inline void WriteTag(TKafkaWritable& writable, TKafkaInt16 version, const typename Meta::Type& value) { + if constexpr (!VersionNone()) { + if (VersionCheck(version) && + VersionCheck(version)) { + if (!IsDefaultValue(value)) { + writable.writeUnsignedVarint(Meta::Tag); + writable.writeUnsignedVarint(TypeStrategy::DoSize(version, value)); + TypeStrategy::DoWriteTag(writable, version, value); + } + } + } +} + +template +inline void ReadTag(TKafkaReadable& readable, TKafkaInt16 version, typename Meta::Type& value) { + if constexpr (!VersionNone()) { + TypeStrategy::DoRead(readable, version, value); + } +} + + +} // NPrivate + +} // namespace NKafka diff --git a/src/library/kafka/kafka_records.cpp b/src/library/kafka/kafka_records.cpp new file mode 100644 index 00000000000..547eccb477d --- /dev/null +++ b/src/library/kafka/kafka_records.cpp @@ -0,0 +1,1045 @@ +#include "kafka_messages_int.h" + +#include + +#include +#include + +#include +#include +#include + +namespace NKafka { + +namespace { + +static constexpr size_t WriteBufferChunkSize = 1 << 16; +static constexpr size_t RecordBatchCrcOffset = + sizeof(TKafkaRecordBatch::BaseOffsetMeta::Type) + + sizeof(TKafkaRecordBatch::BatchLengthMeta::Type) + + sizeof(TKafkaRecordBatch::PartitionLeaderEpochMeta::Type) + + sizeof(TKafkaRecordBatch::MagicMeta::Type); +static constexpr size_t RecordBatchCrcBodyOffset = + RecordBatchCrcOffset + sizeof(TKafkaRecordBatch::CrcMeta::Type); + +void EnsureSupportedCompressionType(ECompressionType compressionType) { + switch (compressionType) { + case ECompressionType::NONE: + case ECompressionType::GZIP: + case ECompressionType::ZSTD: + return; + default: + ythrow yexception() << "unsupported Kafka record batch compression type: " << static_cast(compressionType); + } +} + +void WriteKafkaInt32(TString& data, size_t offset, TKafkaInt32 value) { + data[offset] = static_cast((static_cast(value) >> 24) & 0xff); + data[offset + 1] = static_cast((static_cast(value) >> 16) & 0xff); + data[offset + 2] = static_cast((static_cast(value) >> 8) & 0xff); + data[offset + 3] = static_cast(static_cast(value) & 0xff); +} + +constexpr std::array MakeCrc32cTable() { + constexpr ui32 polynomial = 0x82f63b78; + std::array table = {}; + for (ui32 i = 0; i < table.size(); ++i) { + ui32 crc = i; + for (size_t bit = 0; bit < 8; ++bit) { + crc = (crc >> 1) ^ (polynomial & (0 - (crc & 1))); + } + table[i] = crc; + } + return table; +} + +ui32 Crc32c(const void* data, size_t size) noexcept { + static constexpr auto table = MakeCrc32cTable(); + ui32 crc = ~ui32{0}; + const auto* bytes = static_cast(data); + for (size_t i = 0; i < size; ++i) { + crc = table[(crc ^ bytes[i]) & 0xff] ^ (crc >> 8); + } + return ~crc; +} + +TString DecompressRecordBatchPayload(TStringBuf data, ECompressionType compressionType) { + EnsureSupportedCompressionType(compressionType); + if (compressionType == ECompressionType::NONE) { + return TString(data); + } + + TMemoryInput input(data.data(), data.size()); + switch (compressionType) { + case ECompressionType::GZIP: { + TZLibDecompress gzip(&input, ZLib::GZip); + return gzip.ReadAll(); + } + case ECompressionType::ZSTD: { + TZstdDecompress zstd(&input); + return zstd.ReadAll(); + } + default: + ythrow yexception() << "unsupported Kafka record batch compression type: " << static_cast(compressionType); + } +} + +TString CompressRecordBatchPayload(TStringBuf data, ECompressionType compressionType) { + EnsureSupportedCompressionType(compressionType); + if (compressionType == ECompressionType::NONE) { + return TString(data); + } + + TString result; + TStringOutput output(result); + switch (compressionType) { + case ECompressionType::GZIP: { + TZLibCompress gzip(&output, ZLib::GZip); + gzip.Write(data.data(), data.size()); + gzip.Finish(); + output.Finish(); + return result; + } + case ECompressionType::ZSTD: { + TZstdCompress zstd(&output); + zstd.Write(data.data(), data.size()); + zstd.Finish(); + output.Finish(); + return result; + } + default: + ythrow yexception() << "unsupported Kafka record batch compression type: " << static_cast(compressionType); + } +} + +void EnsureValidRecordBatchRecordsCount(TKafkaInt32 recordsCount) { + if (recordsCount < 0) { + ythrow yexception() << "non-nullable field records was serialized as null"; + } +} + +TString SerializeRecordBatchRecords( + const TKafkaRecordBatch::RecordsMeta::Type& records, + TKafkaVersion version, + bool includeArraySize) { + using RecordsMeta = TKafkaRecordBatch::RecordsMeta; + using ItemStrategy = NPrivate::TypeStrategy< + RecordsMeta, + RecordsMeta::ItemType, + RecordsMeta::ItemTypeDesc>; + + TKafkaWriteBuffer buffer(WriteBufferChunkSize); + TKafkaWritable writable(buffer); + if (includeArraySize) { + NPrivate::TWriteCollector writeCollector; + NPrivate::Write(writeCollector, writable, version, records); + } else { + for (const auto& record : records) { + ItemStrategy::DoWrite(writable, version, record); + } + } + + TString result; + for (auto it = buffer.GetBuffersDeque().rbegin(); it != buffer.GetBuffersDeque().rend(); ++it) { + result.append(it->Data(), it->Size()); + } + return result; +} + +// The returned entries keep TKafkaBytes views into `buffer`, so the caller must +// keep `buffer` alive for as long as the entries are used. +std::vector ReadLegacyRecordEntries(const TBuffer& buffer, TKafkaVersion magic) { + TKafkaReadable readable(buffer); + std::vector entries; + + while (readable.left() > 0) { + const TKafkaVersion entryMagic = readable.take(16); + if (entryMagic != magic) { + ythrow yexception() << "compressed Kafka legacy record magic " << entryMagic + << " does not match wrapper magic " << magic; + } + + auto& entry = entries.emplace_back(); + entry.Read(readable, magic); + } + + return entries; +} + +void AppendLegacyRecord( + TKafkaRecordBatch& batch, + const TKafkaRecordBatchV0& entry, + i64 offset, + std::optional timestamp = std::nullopt) +{ + auto& record = batch.Records.emplace_back(); + record.Length = entry.Record.MessageSize; + record.OffsetDelta = offset; + record.TimestampDelta = timestamp.value_or(entry.Record.Timestamp); + if (entry.Record.Key) { + record.SetKey(TString(entry.Record.Key->data(), entry.Record.Key->size())); + } + if (entry.Record.Value) { + record.SetValue(TString(entry.Record.Value->data(), entry.Record.Value->size())); + } +} + +void AppendLegacyRecords( + TKafkaRecordBatch& batch, + const std::vector& entries, + TKafkaVersion magic, + i64 wrapperOffset = 0, + std::optional wrapperTimestamp = std::nullopt) +{ + if (entries.empty()) { + return; + } + + i64 absoluteBaseOffset = 0; + if (magic == 1) { + absoluteBaseOffset = wrapperOffset == 0 ? 0 : wrapperOffset - entries.back().Offset; + } + + for (const auto& entry : entries) { + const i64 offset = magic == 1 ? absoluteBaseOffset + entry.Offset : entry.Offset; + AppendLegacyRecord(batch, entry, offset, wrapperTimestamp); + } +} + +template +void ForEachLegacyRecordBatch( + TKafkaReadable& recordsReadable, + TKafkaVersion magic, + bool allowCompressed, + TConsumer& consumer) +{ + while (recordsReadable.left() > 0) { + const TKafkaVersion entryMagic = recordsReadable.take(16); + if (entryMagic != magic) { + ythrow yexception() << "Kafka legacy record magic " << entryMagic + << " does not match expected magic " << magic; + } + + TKafkaRecordBatchV0 entry; + entry.Read(recordsReadable, magic); + + const ECompressionType compressionType = entry.Record.CompressionType(); + if (compressionType == ECompressionType::NONE) { + consumer.OnUncompressed(entry); + continue; + } + + if (!allowCompressed) { + ythrow yexception() << "Supported only CompressionType::NONE"; + } + EnsureSupportedCompressionType(compressionType); + if (!entry.Record.Value) { + ythrow yexception() << "compressed Kafka legacy record has null value"; + } + + const auto& value = *entry.Record.Value; + const TString decompressed = DecompressRecordBatchPayload( + TStringBuf(value.data(), value.size()), + compressionType); + TBuffer innerBuffer(decompressed.data(), decompressed.size()); + const std::vector innerEntries = ReadLegacyRecordEntries(innerBuffer, magic); + consumer.OnCompressed(entry, compressionType, innerEntries); + } +} + +struct TLegacyBatchConsumer { + TKafkaRecordBatch& Batch; + TKafkaVersion Magic; + + void OnUncompressed(const TKafkaRecordBatchV0& entry) { + AppendLegacyRecord(Batch, entry, entry.Offset); + } + + void OnCompressed( + const TKafkaRecordBatchV0& entry, + ECompressionType compressionType, + const std::vector& innerEntries) + { + Batch.Attributes = static_cast(compressionType); + const std::optional wrapperTimestamp = Magic == 1 + ? std::optional(entry.Record.Timestamp) + : std::nullopt; + AppendLegacyRecords(Batch, innerEntries, Magic, entry.Offset, wrapperTimestamp); + } +}; + +struct TLegacyHeaderConsumer { + TKafkaBatchHeader& Header; + TKafkaVersion Magic; + + void OnUncompressed(const TKafkaRecordBatchV0& entry) { + AppendHeaderRecord( + entry.Offset, + Magic >= 1 ? entry.Record.Timestamp : 0); + } + + void OnCompressed( + const TKafkaRecordBatchV0& entry, + ECompressionType /*compressionType*/, + const std::vector& innerEntries) + { + Header.Attributes = entry.Record.Attributes; + const std::optional wrapperTimestamp = Magic == 1 + ? std::optional(entry.Record.Timestamp) + : std::nullopt; + + i64 absoluteBaseOffset = 0; + if (Magic == 1) { + absoluteBaseOffset = entry.Offset == 0 ? 0 : entry.Offset - innerEntries.back().Offset; + } + + for (const auto& innerEntry : innerEntries) { + const i64 offset = Magic == 1 + ? absoluteBaseOffset + innerEntry.Offset + : innerEntry.Offset; + const i64 timestamp = wrapperTimestamp.value_or( + Magic >= 1 ? innerEntry.Record.Timestamp : 0); + AppendHeaderRecord(offset, timestamp); + } + } + +private: + void AppendHeaderRecord(i64 offset, i64 timestamp) { + if (Header.RecordsCount == 0) { + Header.BaseOffset = offset; + Header.LastOffsetDelta = 0; + if (Magic >= 1) { + Header.BaseTimestamp = timestamp; + Header.MaxTimestamp = timestamp; + } + } else { + Header.LastOffsetDelta = offset - Header.BaseOffset; + if (Magic >= 1) { + Header.MaxTimestamp = Max(Header.MaxTimestamp, timestamp); + } + } + + ++Header.RecordsCount; + } +}; + +} // namespace + +// +// TKafkaHeader +// +const TKafkaHeader::KeyMeta::Type TKafkaHeader::KeyMeta::Default = std::nullopt; + +TKafkaHeader::TKafkaHeader() + : Key(KeyMeta::Default) { +} + +void TKafkaHeader::Read(TKafkaReadable& _readable, TKafkaVersion _version) { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't read version " << _version << " of TKafkaHeader"; + } + NPrivate::Read(_readable, _version, Key); + NPrivate::Read(_readable, _version, Value); +} + +void TKafkaHeader::Write(TKafkaWritable& _writable, TKafkaVersion _version) const { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't write version " << _version << " of TKafkaHeader"; + } + NPrivate::TWriteCollector _collector; + NPrivate::Write(_collector, _writable, _version, Key); + NPrivate::Write(_collector, _writable, _version, Value); +} + +i32 TKafkaHeader::Size(TKafkaVersion _version) const { + NPrivate::TSizeCollector _collector; + NPrivate::Size(_collector, _version, Key); + NPrivate::Size(_collector, _version, Value); + + return _collector.Size; +} + + +// +// TKafkaRecord +// +const TKafkaRecord::KeyMeta::Type TKafkaRecord::KeyMeta::Default = std::nullopt; + +TKafkaRecord::TKafkaRecord() + : Length(LengthMeta::Default) + , Attributes(AttributesMeta::Default) + , TimestampDelta(TimestampDeltaMeta::Default) + , OffsetDelta(OffsetDeltaMeta::Default) + , Key(KeyMeta::Default) +{ +} + +TKafkaRecord::TKafkaRecord(const TKafkaRecord& other) + : Length(other.Length) + , Attributes(other.Attributes) + , TimestampDelta(other.TimestampDelta) + , OffsetDelta(other.OffsetDelta) + , Key(other.Key) + , Value(other.Value) + , Headers(other.Headers) + , Storage_(other.Storage_) +{ + RebindStorage(); +} + +TKafkaRecord::TKafkaRecord(TKafkaRecord&& other) noexcept + : Length(other.Length) + , Attributes(other.Attributes) + , TimestampDelta(other.TimestampDelta) + , OffsetDelta(other.OffsetDelta) + , Key(other.Key) + , Value(other.Value) + , Headers(std::move(other.Headers)) + , Storage_(std::move(other.Storage_)) +{ + RebindStorage(); +} + +TKafkaRecord& TKafkaRecord::operator=(const TKafkaRecord& other) { + if (this != &other) { + Length = other.Length; + Attributes = other.Attributes; + TimestampDelta = other.TimestampDelta; + OffsetDelta = other.OffsetDelta; + Key = other.Key; + Value = other.Value; + Headers = other.Headers; + Storage_ = other.Storage_; + RebindStorage(); + } + return *this; +} + +TKafkaRecord& TKafkaRecord::operator=(TKafkaRecord&& other) noexcept { + if (this != &other) { + Length = other.Length; + Attributes = other.Attributes; + TimestampDelta = other.TimestampDelta; + OffsetDelta = other.OffsetDelta; + Key = other.Key; + Value = other.Value; + Headers = std::move(other.Headers); + Storage_ = std::move(other.Storage_); + RebindStorage(); + } + return *this; +} + +void TKafkaRecord::SetKey(TString key) { + Storage_.Key = std::move(key); + Key = TArrayRef(Storage_.Key->data(), Storage_.Key->size()); +} + +void TKafkaRecord::SetValue(TString value) { + Storage_.Value = std::move(value); + Value = TArrayRef(Storage_.Value->data(), Storage_.Value->size()); +} + +void TKafkaRecord::AddHeader(TString key, TString value) { + Storage_.Headers.push_back(TStorage::THeaderData{ + .Key = std::move(key), + .Value = std::move(value), + }); + RebindStorage(); +} + +void TKafkaRecord::OwnPayload() { + if (Key) { + Storage_.Key = TString(Key->data(), Key->size()); + } else { + Storage_.Key.reset(); + } + + if (Value) { + Storage_.Value = TString(Value->data(), Value->size()); + } else { + Storage_.Value.reset(); + } + + if (!Headers.empty()) { + std::vector headersStorage; + headersStorage.reserve(Headers.size()); + for (const auto& header : Headers) { + TStorage::THeaderData headerStorage; + if (header.Key) { + headerStorage.Key = TString(header.Key->data(), header.Key->size()); + } + if (header.Value) { + headerStorage.Value = TString(header.Value->data(), header.Value->size()); + } + headersStorage.push_back(std::move(headerStorage)); + } + Storage_.Headers = std::move(headersStorage); + } else { + Storage_.Headers.clear(); + } + + RebindStorage(); +} + +void TKafkaRecord::RebindStorage() { + if (Storage_.Key) { + Key = TArrayRef(Storage_.Key->data(), Storage_.Key->size()); + } + if (Storage_.Value) { + Value = TArrayRef(Storage_.Value->data(), Storage_.Value->size()); + } + if (!Storage_.Headers.empty()) { + Headers.clear(); + Headers.reserve(Storage_.Headers.size()); + for (const auto& headerStorage : Storage_.Headers) { + TKafkaHeader header; + if (headerStorage.Key) { + header.Key = TArrayRef(headerStorage.Key->data(), headerStorage.Key->size()); + } + if (headerStorage.Value) { + header.Value = TArrayRef(headerStorage.Value->data(), headerStorage.Value->size()); + } + Headers.push_back(std::move(header)); + } + } +} + +bool TKafkaRecord::operator==(const TKafkaRecord& other) const { + return Length == other.Length + && Attributes == other.Attributes + && TimestampDelta == other.TimestampDelta + && OffsetDelta == other.OffsetDelta + && Key == other.Key + && Value == other.Value + && Headers == other.Headers; +} + +void TKafkaRecord::Read(TKafkaReadable& _readable, TKafkaVersion _version) { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't read version " << _version << " of TKafkaRecord"; + } + NPrivate::Read(_readable, _version, Length); + NPrivate::Read(_readable, _version, Attributes); + NPrivate::Read(_readable, _version, TimestampDelta); + NPrivate::Read(_readable, _version, OffsetDelta); + NPrivate::Read(_readable, _version, Key); + NPrivate::Read(_readable, _version, Value); + NPrivate::Read(_readable, _version, Headers); +} + +void TKafkaRecord::Write(TKafkaWritable& _writable, TKafkaVersion _version) const { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't write version " << _version << " of TKafkaRecord"; + } + NPrivate::TWriteCollector _collector; + NPrivate::Write(_collector, _writable, _version, Length); + NPrivate::Write(_collector, _writable, _version, Attributes); + NPrivate::Write(_collector, _writable, _version, TimestampDelta); + NPrivate::Write(_collector, _writable, _version, OffsetDelta); + NPrivate::Write(_collector, _writable, _version, Key); + NPrivate::Write(_collector, _writable, _version, Value); + NPrivate::Write(_collector, _writable, _version, Headers); +} + +i32 TKafkaRecord::Size(TKafkaVersion _version) const { + NPrivate::TSizeCollector _collector; + NPrivate::Size(_collector, _version, Length); + NPrivate::Size(_collector, _version, Attributes); + NPrivate::Size(_collector, _version, TimestampDelta); + NPrivate::Size(_collector, _version, OffsetDelta); + NPrivate::Size(_collector, _version, Key); + NPrivate::Size(_collector, _version, Value); + NPrivate::Size(_collector, _version, Headers); + + return _collector.Size; +} + + + +// +// TKafkaRecordBatch +// +TKafkaRecordBatch::TKafkaRecordBatch() + : BaseOffset(BaseOffsetMeta::Default) + , BatchLength(BatchLengthMeta::Default) + , PartitionLeaderEpoch(PartitionLeaderEpochMeta::Default) + , Magic(MagicMeta::Default) + , Crc(CrcMeta::Default) + , Attributes(AttributesMeta::Default) + , LastOffsetDelta(LastOffsetDeltaMeta::Default) + , BaseTimestamp(BaseTimestampMeta::Default) + , MaxTimestamp(MaxTimestampMeta::Default) + , ProducerId(ProducerIdMeta::Default) + , ProducerEpoch(ProducerEpochMeta::Default) + , BaseSequence(BaseSequenceMeta::Default) { +} + +ECompressionType TKafkaRecordBatch::CompressionType() const { + return static_cast(Attributes & 0x07); +} + +ETimestampType TKafkaRecordBatch::TimestampType() { + return (Attributes & 0x08) ? ETimestampType::LOG_APPEND_TIME : ETimestampType::CREATE_TIME; +} + +bool TKafkaRecordBatch::Transactional() { + return Attributes & 0x10; +} + +bool TKafkaRecordBatch::ControlBatch() { + return Attributes & 0x20; +} + +bool TKafkaRecordBatch::HasDeleteHorizonMs() { + return Attributes & 0x40; +} + +void TKafkaRecordBatch::Compress(TKafkaVersion version) { + Y_UNUSED(version); + const auto compressionType = CompressionType(); + EnsureSupportedCompressionType(compressionType); +} + +void TKafkaRecordBatch::Decompress(TKafkaVersion version) { + Y_UNUSED(version); + EnsureSupportedCompressionType(CompressionType()); +} + +void TKafkaRecordBatch::Read(TKafkaReadable& _readable, TKafkaVersion _version) { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't read version " << _version << " of TKafkaRecordBatch"; + } + NPrivate::Read(_readable, _version, BaseOffset); + NPrivate::Read(_readable, _version, BatchLength); + NPrivate::Read(_readable, _version, PartitionLeaderEpoch); + NPrivate::Read(_readable, _version, Magic); + if (2 != Magic) { + ythrow yexception() << "Supported only RecordBatch version 2 but " << (ui16)Magic; + } + + NPrivate::Read(_readable, _version, Crc); + NPrivate::Read(_readable, _version, Attributes); + if (CompressionType() != ECompressionType::NONE && !_readable.GetAllowCompressed()) { + ythrow yexception() << "Supported only CompressionType::NONE"; + } + EnsureSupportedCompressionType(CompressionType()); + + NPrivate::Read(_readable, _version, LastOffsetDelta); + NPrivate::Read(_readable, _version, BaseTimestamp); + NPrivate::Read(_readable, _version, MaxTimestamp); + NPrivate::Read(_readable, _version, ProducerId); + NPrivate::Read(_readable, _version, ProducerEpoch); + NPrivate::Read(_readable, _version, BaseSequence); + RecordsMeta::Type().swap(Records); + + if (CompressionType() == ECompressionType::NONE) { + NPrivate::Read(_readable, _version, Records); + for (auto& record : Records) { + record.OwnPayload(); + } + } else { + const TKafkaInt32 recordsCount = NPrivate::ReadArraySize(_readable, _version); + EnsureValidRecordBatchRecordsCount(recordsCount); + const auto compressed = _readable.Bytes(_readable.left()); + const TString decompressed = DecompressRecordBatchPayload( + TStringBuf(compressed.data(), compressed.size()), CompressionType()); + TBuffer buffer(decompressed.data(), decompressed.size()); + TKafkaReadable recordsReadable(buffer); + Records.resize(recordsCount); + using ItemStrategy = NPrivate::TypeStrategy< + RecordsMeta, + RecordsMeta::ItemType, + RecordsMeta::ItemTypeDesc>; + for (auto& record : Records) { + ItemStrategy::DoRead(recordsReadable, _version, record); + } + for (auto& record : Records) { + record.OwnPayload(); + } + } +} + +void TKafkaRecordBatch::Write(TKafkaWritable& _writable, TKafkaVersion _version) const { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't write version " << _version << " of TKafkaRecordBatch"; + } + EnsureSupportedCompressionType(CompressionType()); + NPrivate::TWriteCollector _collector; + NPrivate::Write(_collector, _writable, _version, BaseOffset); + NPrivate::Write(_collector, _writable, _version, BatchLength); + NPrivate::Write(_collector, _writable, _version, PartitionLeaderEpoch); + NPrivate::Write(_collector, _writable, _version, Magic); + NPrivate::Write(_collector, _writable, _version, Crc); + NPrivate::Write(_collector, _writable, _version, Attributes); + NPrivate::Write(_collector, _writable, _version, LastOffsetDelta); + NPrivate::Write(_collector, _writable, _version, BaseTimestamp); + NPrivate::Write(_collector, _writable, _version, MaxTimestamp); + NPrivate::Write(_collector, _writable, _version, ProducerId); + NPrivate::Write(_collector, _writable, _version, ProducerEpoch); + NPrivate::Write(_collector, _writable, _version, BaseSequence); + if (CompressionType() == ECompressionType::NONE) { + NPrivate::Write(_collector, _writable, _version, Records); + } else { + const TString packedRecords = CompressRecordBatchPayload( + SerializeRecordBatchRecords(Records, _version, false), CompressionType()); + NPrivate::WriteArraySize(_writable, _version, Records.size()); + _writable.write(packedRecords.data(), packedRecords.size()); + } +} + +i32 TKafkaRecordBatch::Size(TKafkaVersion _version) const { + EnsureSupportedCompressionType(CompressionType()); + NPrivate::TSizeCollector _collector; + NPrivate::Size(_collector, _version, BaseOffset); + NPrivate::Size(_collector, _version, BatchLength); + NPrivate::Size(_collector, _version, PartitionLeaderEpoch); + NPrivate::Size(_collector, _version, Magic); + NPrivate::Size(_collector, _version, Crc); + NPrivate::Size(_collector, _version, Attributes); + NPrivate::Size(_collector, _version, LastOffsetDelta); + NPrivate::Size(_collector, _version, BaseTimestamp); + NPrivate::Size(_collector, _version, MaxTimestamp); + NPrivate::Size(_collector, _version, ProducerId); + NPrivate::Size(_collector, _version, ProducerEpoch); + NPrivate::Size(_collector, _version, BaseSequence); + if (CompressionType() == ECompressionType::NONE) { + NPrivate::Size(_collector, _version, Records); + } else { + const TString packedRecords = CompressRecordBatchPayload( + SerializeRecordBatchRecords(Records, _version, false), CompressionType()); + _collector.Size += NPrivate::ArraySize(_version, Records.size()) + packedRecords.size(); + } + + return _collector.Size; +} + + +// +// TKafkaBatchHeader +// +TKafkaBatchHeader::TKafkaBatchHeader() + : BaseOffset(BaseOffsetMeta::Default) + , BatchLength(BatchLengthMeta::Default) + , PartitionLeaderEpoch(PartitionLeaderEpochMeta::Default) + , Magic(MagicMeta::Default) + , Crc(CrcMeta::Default) + , Attributes(AttributesMeta::Default) + , LastOffsetDelta(LastOffsetDeltaMeta::Default) + , BaseTimestamp(BaseTimestampMeta::Default) + , MaxTimestamp(MaxTimestampMeta::Default) + , ProducerId(ProducerIdMeta::Default) + , ProducerEpoch(ProducerEpochMeta::Default) + , BaseSequence(BaseSequenceMeta::Default) + , RecordsCount(RecordsCountMeta::Default) { +} + +void TKafkaBatchHeader::Read(TKafkaReadable& _readable, TKafkaVersion _version) { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't read version " << _version << " of TKafkaBatchHeader"; + } + NPrivate::Read(_readable, _version, BaseOffset); + NPrivate::Read(_readable, _version, BatchLength); + NPrivate::Read(_readable, _version, PartitionLeaderEpoch); + NPrivate::Read(_readable, _version, Magic); + NPrivate::Read(_readable, _version, Crc); + NPrivate::Read(_readable, _version, Attributes); + NPrivate::Read(_readable, _version, LastOffsetDelta); + NPrivate::Read(_readable, _version, BaseTimestamp); + NPrivate::Read(_readable, _version, MaxTimestamp); + NPrivate::Read(_readable, _version, ProducerId); + NPrivate::Read(_readable, _version, ProducerEpoch); + NPrivate::Read(_readable, _version, BaseSequence); + NPrivate::Read(_readable, _version, RecordsCount); + EnsureValidRecordBatchRecordsCount(RecordsCount); +} + +void TKafkaBatchHeader::Write(TKafkaWritable& _writable, TKafkaVersion _version) const { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't write version " << _version << " of TKafkaBatchHeader"; + } + NPrivate::TWriteCollector _collector; + NPrivate::Write(_collector, _writable, _version, BaseOffset); + NPrivate::Write(_collector, _writable, _version, BatchLength); + NPrivate::Write(_collector, _writable, _version, PartitionLeaderEpoch); + NPrivate::Write(_collector, _writable, _version, Magic); + NPrivate::Write(_collector, _writable, _version, Crc); + NPrivate::Write(_collector, _writable, _version, Attributes); + NPrivate::Write(_collector, _writable, _version, LastOffsetDelta); + NPrivate::Write(_collector, _writable, _version, BaseTimestamp); + NPrivate::Write(_collector, _writable, _version, MaxTimestamp); + NPrivate::Write(_collector, _writable, _version, ProducerId); + NPrivate::Write(_collector, _writable, _version, ProducerEpoch); + NPrivate::Write(_collector, _writable, _version, BaseSequence); + NPrivate::Write(_collector, _writable, _version, RecordsCount); +} + +i32 TKafkaBatchHeader::Size(TKafkaVersion _version) const { + NPrivate::TSizeCollector _collector; + NPrivate::Size(_collector, _version, BaseOffset); + NPrivate::Size(_collector, _version, BatchLength); + NPrivate::Size(_collector, _version, PartitionLeaderEpoch); + NPrivate::Size(_collector, _version, Magic); + NPrivate::Size(_collector, _version, Crc); + NPrivate::Size(_collector, _version, Attributes); + NPrivate::Size(_collector, _version, LastOffsetDelta); + NPrivate::Size(_collector, _version, BaseTimestamp); + NPrivate::Size(_collector, _version, MaxTimestamp); + NPrivate::Size(_collector, _version, ProducerId); + NPrivate::Size(_collector, _version, ProducerEpoch); + NPrivate::Size(_collector, _version, BaseSequence); + NPrivate::Size(_collector, _version, RecordsCount); + return _collector.Size; +} + + + + +// +// TKafkaRecordV0 +// +const TKafkaRecordV0::KeyMeta::Type TKafkaRecordV0::KeyMeta::Default = std::nullopt; + +ECompressionType TKafkaRecordV0::CompressionType() const { + return static_cast(Attributes & 0x07); +} + +TKafkaRecordV0::TKafkaRecordV0() + : MessageSize(MessageSizeMeta::Default) + , Crc(CrcMeta::Default) + , Magic(MagicMeta::Default) + , Attributes(AttributesMeta::Default) + , Timestamp(TimestampMeta::Default) + , Key(KeyMeta::Default) { +} + +void TKafkaRecordV0::Read(TKafkaReadable& _readable, TKafkaVersion _version) { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't read version " << _version << " of TKafkaRecordV0"; + } + NPrivate::Read(_readable, _version, MessageSize); + NPrivate::Read(_readable, _version, Crc); + NPrivate::Read(_readable, _version, Magic); + NPrivate::Read(_readable, _version, Attributes); + NPrivate::Read(_readable, _version, Timestamp); + NPrivate::Read(_readable, _version, Key); + NPrivate::Read(_readable, _version, Value); +} + +void TKafkaRecordV0::Write(TKafkaWritable& _writable, TKafkaVersion _version) const { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't write version " << _version << " of TKafkaRecordV0"; + } + NPrivate::TWriteCollector _collector; + NPrivate::Write(_collector, _writable, _version, MessageSize); + NPrivate::Write(_collector, _writable, _version, Crc); + NPrivate::Write(_collector, _writable, _version, Magic); + NPrivate::Write(_collector, _writable, _version, Attributes); + NPrivate::Write(_collector, _writable, _version, Timestamp); + NPrivate::Write(_collector, _writable, _version, Key); + NPrivate::Write(_collector, _writable, _version, Value); +} + +i32 TKafkaRecordV0::Size(TKafkaVersion _version) const { + NPrivate::TSizeCollector _collector; + NPrivate::Size(_collector, _version, MessageSize); + NPrivate::Size(_collector, _version, Crc); + NPrivate::Size(_collector, _version, Magic); + NPrivate::Size(_collector, _version, Attributes); + NPrivate::Size(_collector, _version, Timestamp); + NPrivate::Size(_collector, _version, Key); + NPrivate::Size(_collector, _version, Value); + + return _collector.Size; +} + + +// +// TKafkaRecordBatchV0 +// +TKafkaRecordBatchV0::TKafkaRecordBatchV0() + : Offset(OffsetMeta::Default) { +} + +void TKafkaRecordBatchV0::Read(TKafkaReadable& _readable, TKafkaVersion _version) { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't read version " << _version << " of TKafkaRecordBatchV0"; + } + NPrivate::Read(_readable, _version, Offset); + NPrivate::Read(_readable, _version, Record); +} + +void TKafkaRecordBatchV0::Write(TKafkaWritable& _writable, TKafkaVersion _version) const { + if (!NPrivate::VersionCheck(_version)) { + ythrow yexception() << "Can't write version " << _version << " of TKafkaRecordBatchV0"; + } + NPrivate::TWriteCollector _collector; + NPrivate::Write(_collector, _writable, _version, Offset); + NPrivate::Write(_collector, _writable, _version, Record); +} + +i32 TKafkaRecordBatchV0::Size(TKafkaVersion _version) const { + NPrivate::TSizeCollector _collector; + NPrivate::Size(_collector, _version, Offset); + NPrivate::Size(_collector, _version, Record); + + return _collector.Size; +} + +void NPrivate::ReadLegacyRecordBatch( + TKafkaReadable& readable, + TKafkaVersion magic, + size_t length, + TKafkaRecordBatch& batch) +{ + const auto data = readable.Bytes(length); + TBuffer buffer(data.data(), data.size()); + TKafkaReadable recordsReadable(buffer); + + batch = {}; + batch.Magic = 2; + batch.BaseOffset = 0; + batch.BaseTimestamp = 0; + + TLegacyBatchConsumer consumer{batch, magic}; + ForEachLegacyRecordBatch(recordsReadable, magic, readable.GetAllowCompressed(), consumer); + + batch.LastOffsetDelta = batch.Records.empty() ? 0 : batch.Records.back().OffsetDelta; +} + +void ReadLegacyRecordBatchHeaderFields( + TKafkaReadable& recordsReadable, + TKafkaVersion magic, + bool allowCompressed, + TKafkaBatchHeader& header) +{ + TLegacyHeaderConsumer consumer{header, magic}; + ForEachLegacyRecordBatch(recordsReadable, magic, allowCompressed, consumer); +} + +TKafkaBatchHeader ReadLegacyRecordBatchHeader( + TKafkaReadable& readable, + TKafkaVersion magic, + size_t length) +{ + const auto data = readable.Bytes(length); + TBuffer buffer(data.data(), data.size()); + TKafkaReadable recordsReadable(buffer); + + // Legacy batches (magic 0/1) do not have a v2 RecordBatch header on wire. + // BaseSequence, ProducerId, Crc, etc. stay at ctor defaults. + TKafkaBatchHeader header; + header.Magic = magic; + ReadLegacyRecordBatchHeaderFields(recordsReadable, magic, readable.GetAllowCompressed(), header); + return header; +} + +TKafkaRecordBatch ReadKafkaRecordBatch(TStringBuf data, TKafkaVersion version) { + TBuffer buffer(data.data(), data.size()); + TKafkaReadable readable(buffer); + readable.SetAllowCompressed(true); + + TKafkaRecordBatch batch; + batch.Read(readable, version); + if (readable.left() != 0) { + ythrow yexception() << "unexpected extra bytes after Kafka record batch: " << readable.left(); + } + return batch; +} + +TKafkaRecordBatch ReadRecordBatch(TStringBuf data) { + static constexpr size_t RecordBatchMagicOffset = + sizeof(TKafkaInt64) + sizeof(TKafkaInt32) + sizeof(TKafkaInt32); + + if (data.size() <= RecordBatchMagicOffset) { + ythrow yexception() << "Kafka record batch is too small: " << data.size(); + } + + const auto magic = static_cast(static_cast(data[RecordBatchMagicOffset])); + if (magic >= TKafkaRecordBatch::MagicMeta::Default) { + return ReadKafkaRecordBatch(data); + } + + TBuffer buffer(data.data(), data.size()); + TKafkaReadable readable(buffer); + TKafkaRecordBatch batch; + NPrivate::ReadLegacyRecordBatch(readable, magic, data.size(), batch); + return batch; +} + +TString WriteKafkaRecordBatch(const TKafkaRecordBatch& batch, TKafkaVersion version) { + TKafkaWriteBuffer buffer(WriteBufferChunkSize); + TKafkaWritable writable(buffer); + batch.Write(writable, version); + TString result = buffer.AsString(); + if (result.size() >= RecordBatchCrcBodyOffset) { + const ui32 crc = Crc32c(result.data() + RecordBatchCrcBodyOffset, result.size() - RecordBatchCrcBodyOffset); + WriteKafkaInt32(result, RecordBatchCrcOffset, static_cast(crc)); + } + return result; +} + +std::pair GetBatchBaseSeqNo(const TKafkaBatchHeader& header) { + if (header.ProducerId >= 0) { + if (header.BaseSequence < 0) { + return {EKafkaErrors::INVALID_RECORD, 0}; + } + return {EKafkaErrors::NONE_ERROR, static_cast(header.BaseSequence)}; + } + + if (header.BaseOffset < 0) { + return {EKafkaErrors::INVALID_RECORD, 0}; + } + return {EKafkaErrors::NONE_ERROR, static_cast(header.BaseOffset)}; +} + +std::pair GetBatchMaxSeqNo(const TKafkaBatchHeader& header, ui64 baseSeqNo) { + if (header.ProducerId >= 0) { + return { + EKafkaErrors::NONE_ERROR, + (baseSeqNo + header.RecordsCount - 1) % (static_cast(std::numeric_limits::max()) + 1) + }; + } + + if (header.LastOffsetDelta < 0) { + return {EKafkaErrors::INVALID_RECORD, 0}; + } + return {EKafkaErrors::NONE_ERROR, baseSeqNo + static_cast(header.LastOffsetDelta)}; +} + +std::optional ReadKafkaBatchHeader(TStringBuf data, TKafkaVersion version) { + static constexpr size_t RecordBatchMagicOffset = + sizeof(TKafkaInt64) + sizeof(TKafkaInt32) + sizeof(TKafkaInt32); + + if (data.size() <= RecordBatchMagicOffset) { + return std::nullopt; + } + + const auto magic = static_cast(static_cast(data[RecordBatchMagicOffset])); + + TBuffer buffer(data.data(), data.size()); + TKafkaReadable readable(buffer); + readable.SetAllowCompressed(true); + + try { + if (magic < TKafkaRecordBatch::MagicMeta::Default) { + return ReadLegacyRecordBatchHeader(readable, magic, data.size()); + } + + TKafkaBatchHeader header; + header.Read(readable, version); + return header; + } catch (const yexception&) { + return std::nullopt; + } +} + +ui64 GetRecordSeqNo(const TKafkaRecordBatch& batch, size_t recordIndex, const TKafkaRecord& record) { + if (batch.ProducerId >= 0) { + return (static_cast(batch.BaseSequence) + recordIndex) + % (static_cast(std::numeric_limits::max()) + 1); + } + return static_cast(batch.BaseOffset) + record.OffsetDelta; +} + +} // namespace NKafka diff --git a/src/library/kafka/kafka_records.h b/src/library/kafka/kafka_records.h new file mode 100644 index 00000000000..5b5fb44c603 --- /dev/null +++ b/src/library/kafka/kafka_records.h @@ -0,0 +1,690 @@ +#pragma once + +#include "kafka.h" + +#include +#include + +namespace NKafka { + +enum ETimestampType { + CREATE_TIME = 0, + LOG_APPEND_TIME = 1 +}; + +class TKafkaHeader: public TMessage { +public: + struct MessageMeta { + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + + TKafkaHeader(); + ~TKafkaHeader() = default; + + struct KeyMeta { + using Type = TKafkaBytes; + using TypeDesc = NPrivate::TKafkaBytesDesc; + + static constexpr const char* Name = "key"; + static constexpr const char* About = ""; + static const Type Default; // = {""}; + static constexpr NPrivate::ESizeFormat SizeFormat = NPrivate::ESizeFormat::Varint; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsAlways; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + KeyMeta::Type Key; + + struct ValueMeta { + using Type = TKafkaBytes; + using TypeDesc = NPrivate::TKafkaBytesDesc; + + static constexpr const char* Name = "value"; + static constexpr const char* About = ""; + static constexpr NPrivate::ESizeFormat SizeFormat = NPrivate::ESizeFormat::Varint; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + ValueMeta::Type Value; + + TString CodecKeyStr; + TString CodecValueStr; + + i32 Size(TKafkaVersion version) const override; + void Read(TKafkaReadable& readable, TKafkaVersion version) override; + void Write(TKafkaWritable& writable, TKafkaVersion version) const override; + + bool operator==(const TKafkaHeader& other) const = default; +}; + +class TKafkaRecord: public TMessage { +public: + struct MessageMeta { + static constexpr TKafkaVersions PresentVersions = TKafkaVersions(2, 2); + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + + TKafkaRecord(); + TKafkaRecord(const TKafkaRecord& other); + TKafkaRecord(TKafkaRecord&& other) noexcept; + TKafkaRecord& operator=(const TKafkaRecord& other); + TKafkaRecord& operator=(TKafkaRecord&& other) noexcept; + ~TKafkaRecord() = default; + + void SetKey(TString key); + void SetValue(TString value); + void AddHeader(TString key, TString value); + void OwnPayload(); + + struct LengthMeta { + using Type = TKafkaInt32; + using TypeDesc = NPrivate::TKafkaVarintDesc; + + static constexpr const char* Name = "length"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + LengthMeta::Type Length; + + struct AttributesMeta { + using Type = TKafkaInt8; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "attributes"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + AttributesMeta::Type Attributes; + + struct TimestampDeltaMeta { + using Type = TKafkaInt64; + using TypeDesc = NPrivate::TKafkaVarintDesc; + + static constexpr const char* Name = "timestampDelta"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + TimestampDeltaMeta::Type TimestampDelta; + + struct OffsetDeltaMeta { + using Type = TKafkaInt64; + using TypeDesc = NPrivate::TKafkaVarintDesc; + + static constexpr const char* Name = "offsetDelta"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + OffsetDeltaMeta::Type OffsetDelta; + + struct KeyMeta { + using Type = TKafkaBytes; + using TypeDesc = NPrivate::TKafkaBytesDesc; + + static constexpr const char* Name = "key"; + static constexpr const char* About = ""; + static const Type Default; // = {""}; + static constexpr NPrivate::ESizeFormat SizeFormat = NPrivate::ESizeFormat::Varint; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsAlways; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + KeyMeta::Type Key; + + struct ValueMeta { + using Type = TKafkaBytes; + using TypeDesc = NPrivate::TKafkaBytesDesc; + + static constexpr const char* Name = "value"; + static constexpr const char* About = ""; + static constexpr NPrivate::ESizeFormat SizeFormat = NPrivate::ESizeFormat::Varint; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsAlways; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + ValueMeta::Type Value; + + struct HeadersMeta { + using ItemType = TKafkaHeader; + using ItemTypeDesc = NPrivate::TKafkaStructDesc; + using Type = std::vector; + using TypeDesc = NPrivate::TKafkaArrayDesc; + + static constexpr const char* Name = "headers"; + static constexpr const char* About = ""; + static constexpr NPrivate::ESizeFormat SizeFormat = NPrivate::ESizeFormat::Varint; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsAlways; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + HeadersMeta::Type Headers; + + i32 Size(TKafkaVersion version) const override; + void Read(TKafkaReadable& readable, TKafkaVersion version) override; + void Write(TKafkaWritable& writable, TKafkaVersion version) const override; + + bool operator==(const TKafkaRecord& other) const; + +private: + struct TStorage { + struct THeaderData { + std::optional Key; + std::optional Value; + }; + + std::optional Key; + std::optional Value; + std::vector Headers; + }; + + TStorage Storage_; + + void RebindStorage(); +}; + +namespace NPrivate { + +void ReadLegacyRecordBatch( + TKafkaReadable& readable, + TKafkaVersion magic, + size_t length, + TKafkaRecordBatch& batch); + +} // namespace NPrivate + +class TKafkaRecordBatch: public TMessage { +public: + struct MessageMeta { + static constexpr TKafkaVersions PresentVersions = TKafkaVersions(2, 2); + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + + TKafkaRecordBatch(); + ~TKafkaRecordBatch() = default; + + struct BaseOffsetMeta { + using Type = TKafkaInt64; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "baseOffset"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + BaseOffsetMeta::Type BaseOffset; + + struct BatchLengthMeta { + using Type = TKafkaInt32; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "batchLength"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + BatchLengthMeta::Type BatchLength; + + struct PartitionLeaderEpochMeta { + using Type = TKafkaInt32; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "partitionLeaderEpoch"; + static constexpr const char* About = ""; + static constexpr Type Default = -1; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + PartitionLeaderEpochMeta::Type PartitionLeaderEpoch; + + struct MagicMeta { + using Type = TKafkaInt8; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "magic"; + static constexpr const char* About = "current magic value is 2"; + static constexpr Type Default = 2; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + MagicMeta::Type Magic; + + struct CrcMeta { + using Type = TKafkaInt32; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "crc"; + static constexpr const char* About = "The CRC covers the data from the attributes to the end of the batch (i.e. all the bytes that follow the CRC)"; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + CrcMeta::Type Crc; + + struct AttributesMeta { + using Type = TKafkaInt16; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "attributes"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + AttributesMeta::Type Attributes; + + struct LastOffsetDeltaMeta { + using Type = TKafkaInt32; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "lastOffsetDelta"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + LastOffsetDeltaMeta::Type LastOffsetDelta; + + struct BaseTimestampMeta { + using Type = TKafkaInt64; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "baseTimestamp"; + static constexpr const char* About = ""; + static constexpr Type Default = -1; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + BaseTimestampMeta::Type BaseTimestamp; + + struct MaxTimestampMeta { + using Type = TKafkaInt64; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "maxTimestamp"; + static constexpr const char* About = ""; + static constexpr Type Default = -1; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + MaxTimestampMeta::Type MaxTimestamp; + + struct ProducerIdMeta { + using Type = TKafkaInt64; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "producerId"; + static constexpr const char* About = ""; + static constexpr Type Default = -1; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + ProducerIdMeta::Type ProducerId; + + struct ProducerEpochMeta { + using Type = TKafkaInt16; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "producerEpoch"; + static constexpr const char* About = ""; + static constexpr Type Default = -1; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + ProducerEpochMeta::Type ProducerEpoch; + + struct BaseSequenceMeta { + using Type = TKafkaInt32; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "baseSequence"; + static constexpr const char* About = ""; + static constexpr Type Default = -1; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + BaseSequenceMeta::Type BaseSequence; + + struct RecordsMeta { + using ItemType = TKafkaRecord; + using ItemTypeDesc = NPrivate::TKafkaStructDesc; + using Type = std::vector; + using TypeDesc = NPrivate::TKafkaArrayDesc; + + static constexpr const char* Name = "records"; + static constexpr const char* About = ""; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + RecordsMeta::Type Records; + + i32 Size(TKafkaVersion version) const override; + void Read(TKafkaReadable& readable, TKafkaVersion version) override; + void Write(TKafkaWritable& writable, TKafkaVersion version) const override; + + bool operator==(const TKafkaRecordBatch& other) const = default; + + ECompressionType CompressionType() const; + ETimestampType TimestampType(); + bool Transactional(); + bool ControlBatch(); + bool HasDeleteHorizonMs(); + + void Compress(TKafkaVersion version = MessageMeta::PresentVersions.Max); + void Decompress(TKafkaVersion version = MessageMeta::PresentVersions.Max); +}; + + +struct TKafkaBatchHeader: public TMessage { + using MessageMeta = TKafkaRecordBatch::MessageMeta; + using BaseOffsetMeta = TKafkaRecordBatch::BaseOffsetMeta; + using BatchLengthMeta = TKafkaRecordBatch::BatchLengthMeta; + using PartitionLeaderEpochMeta = TKafkaRecordBatch::PartitionLeaderEpochMeta; + using MagicMeta = TKafkaRecordBatch::MagicMeta; + using CrcMeta = TKafkaRecordBatch::CrcMeta; + using AttributesMeta = TKafkaRecordBatch::AttributesMeta; + using LastOffsetDeltaMeta = TKafkaRecordBatch::LastOffsetDeltaMeta; + using BaseTimestampMeta = TKafkaRecordBatch::BaseTimestampMeta; + using MaxTimestampMeta = TKafkaRecordBatch::MaxTimestampMeta; + using ProducerIdMeta = TKafkaRecordBatch::ProducerIdMeta; + using ProducerEpochMeta = TKafkaRecordBatch::ProducerEpochMeta; + using BaseSequenceMeta = TKafkaRecordBatch::BaseSequenceMeta; + + struct RecordsCountMeta { + using Type = TKafkaInt32; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "recordsCount"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + + TKafkaBatchHeader(); + ~TKafkaBatchHeader() = default; + + BaseOffsetMeta::Type BaseOffset; + BatchLengthMeta::Type BatchLength; + PartitionLeaderEpochMeta::Type PartitionLeaderEpoch; + MagicMeta::Type Magic; + CrcMeta::Type Crc; + AttributesMeta::Type Attributes; + LastOffsetDeltaMeta::Type LastOffsetDelta; + BaseTimestampMeta::Type BaseTimestamp; + MaxTimestampMeta::Type MaxTimestamp; + ProducerIdMeta::Type ProducerId; + ProducerEpochMeta::Type ProducerEpoch; + BaseSequenceMeta::Type BaseSequence; + RecordsCountMeta::Type RecordsCount; + + i32 Size(TKafkaVersion version) const override; + void Read(TKafkaReadable& readable, TKafkaVersion version) override; + void Write(TKafkaWritable& writable, TKafkaVersion version) const override; + + bool operator==(const TKafkaBatchHeader& other) const = default; +}; + +class TKafkaRecordV0: public TMessage { +public: + struct MessageMeta { + static constexpr TKafkaVersions PresentVersions = TKafkaVersions(0, 1); + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + + TKafkaRecordV0(); + ~TKafkaRecordV0() = default; + + struct MessageSizeMeta { + using Type = TKafkaInt32; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "messageSize"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + MessageSizeMeta::Type MessageSize; + + struct CrcMeta { + using Type = TKafkaInt32; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "CRC"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + CrcMeta::Type Crc; + + struct MagicMeta { + using Type = TKafkaInt8; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "magic"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + MagicMeta::Type Magic; + + struct AttributesMeta { + using Type = TKafkaInt8; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "attributes"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + AttributesMeta::Type Attributes; + + struct TimestampMeta { + using Type = TKafkaInt64; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "timestamp"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = TKafkaVersions(1, Max()); + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsAlways; + }; + TimestampMeta::Type Timestamp; + + struct KeyMeta { + using Type = TKafkaBytes; + using TypeDesc = NPrivate::TKafkaBytesDesc; + + static constexpr const char* Name = "key"; + static constexpr const char* About = ""; + static const Type Default; // = {""}; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsAlways; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + KeyMeta::Type Key; + + struct ValueMeta { + using Type = TKafkaBytes; + using TypeDesc = NPrivate::TKafkaBytesDesc; + + static constexpr const char* Name = "value"; + static constexpr const char* About = ""; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + ValueMeta::Type Value; + + i32 Size(TKafkaVersion version) const override; + void Read(TKafkaReadable& readable, TKafkaVersion version) override; + void Write(TKafkaWritable& writable, TKafkaVersion version) const override; + + bool operator==(const TKafkaRecordV0& other) const = default; + + ECompressionType CompressionType() const; +}; + +class TKafkaRecordBatchV0: public TMessage { +public: + struct MessageMeta { + static constexpr TKafkaVersions PresentVersions = TKafkaVersions(0, 1); + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + + TKafkaRecordBatchV0(); + ~TKafkaRecordBatchV0() = default; + + struct OffsetMeta { + using Type = TKafkaInt64; + using TypeDesc = NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "offset"; + static constexpr const char* About = ""; + static constexpr Type Default = 0; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + OffsetMeta::Type Offset; + + struct RecordMeta { + using Type = TKafkaRecordV0; + using TypeDesc = NPrivate::TKafkaStructDesc; + + static constexpr const char* Name = "records"; + static constexpr const char* About = ""; + + static constexpr TKafkaVersions PresentVersions = VersionsAlways; + static constexpr TKafkaVersions TaggedVersions = VersionsNever; + static constexpr TKafkaVersions NullableVersions = VersionsNever; + static constexpr TKafkaVersions FlexibleVersions = VersionsNever; + }; + RecordMeta::Type Record; + + i32 Size(TKafkaVersion version) const override; + void Read(TKafkaReadable& readable, TKafkaVersion version) override; + void Write(TKafkaWritable& writable, TKafkaVersion version) const override; + + bool operator==(const TKafkaRecordBatchV0& other) const = default; +}; + +std::optional ReadKafkaBatchHeader( + TStringBuf data, + TKafkaVersion version = 2); + +TKafkaBatchHeader ReadLegacyRecordBatchHeader( + TKafkaReadable& readable, + TKafkaVersion magic, + size_t length); + +TKafkaRecordBatch ReadKafkaRecordBatch( + TStringBuf data, + TKafkaVersion version = 2); +TKafkaRecordBatch ReadRecordBatch(TStringBuf data); +TString WriteKafkaRecordBatch(const TKafkaRecordBatch& batch, TKafkaVersion version = 2); + +std::pair GetBatchBaseSeqNo(const TKafkaBatchHeader& header); +std::pair GetBatchMaxSeqNo(const TKafkaBatchHeader& header, ui64 baseSeqNo); +ui64 GetRecordSeqNo(const TKafkaRecordBatch& batch, size_t recordIndex, const TKafkaRecord& record); + +} // namespace NKafka diff --git a/src/library/kafka/kafka_write_buffer.cpp b/src/library/kafka/kafka_write_buffer.cpp new file mode 100644 index 00000000000..daeda25f78d --- /dev/null +++ b/src/library/kafka/kafka_write_buffer.cpp @@ -0,0 +1,59 @@ +#include + +#include "kafka_write_buffer.h" + +namespace NKafka { + +TKafkaWriteBuffer::TKafkaWriteBuffer(size_t chunkSize) + : ChunkSize_(chunkSize) +{ +} + +void TKafkaWriteBuffer::write(const char* src, size_t length) { + size_t left = length; + size_t offset = 0; + + do { + TBuffer& buffer = GetOrCreateFrontBuffer(); + if (buffer.Avail() < left) { + const size_t avail = buffer.Avail(); + buffer.Append(src + offset, avail); + offset += avail; + left -= avail; + Buffers_.push_front(TBuffer(ChunkSize_)); + } else { + buffer.Append(src + offset, left); + break; + } + } while (left > 0); +} + +TBuffer& TKafkaWriteBuffer::GetFrontBuffer() { + return GetOrCreateFrontBuffer(); +} + +const TBuffer& TKafkaWriteBuffer::GetFrontBuffer() const { + Y_ABORT_UNLESS(!Buffers_.empty()); + return Buffers_.front(); +} + +const TDeque& TKafkaWriteBuffer::GetBuffersDeque() const { + return Buffers_; +} + +TString TKafkaWriteBuffer::AsString() const { + TString result; + for (auto it = Buffers_.rbegin(); it != Buffers_.rend(); ++it) { + result.append(it->Data(), it->Size()); + } + return result; +} + +TBuffer& TKafkaWriteBuffer::GetOrCreateFrontBuffer() { + if (Buffers_.empty()) { + Buffers_.push_front(TBuffer(ChunkSize_)); + } + return Buffers_.front(); +} + +} // namespace NKafka diff --git a/src/library/kafka/kafka_write_buffer.h b/src/library/kafka/kafka_write_buffer.h new file mode 100644 index 00000000000..e9a7be5df77 --- /dev/null +++ b/src/library/kafka/kafka_write_buffer.h @@ -0,0 +1,28 @@ +#pragma once + +#include +#include +#include + +namespace NKafka { + +class TKafkaWriteBuffer { +public: + explicit TKafkaWriteBuffer(size_t chunkSize); + + void write(const char* src, size_t length); + + TBuffer& GetFrontBuffer(); + const TBuffer& GetFrontBuffer() const; + const TDeque& GetBuffersDeque() const; + + TString AsString() const; + +private: + TBuffer& GetOrCreateFrontBuffer(); + + size_t ChunkSize_; + TDeque Buffers_; +}; + +} // namespace NKafka diff --git a/src/library/kafka/ut/kafka_messages_int_ut.cpp b/src/library/kafka/ut/kafka_messages_int_ut.cpp new file mode 100644 index 00000000000..da3818797d9 --- /dev/null +++ b/src/library/kafka/ut/kafka_messages_int_ut.cpp @@ -0,0 +1,334 @@ +#include + +#include + +namespace NKafka { +namespace { + +static constexpr size_t BUFFER_SIZE = 1 << 16; + +#define SIMPLE_HEAD(Type_, Value) \ + Meta_##Type_::Type value = Value; \ + Meta_##Type_::Type result; \ + \ + TKafkaWriteBuffer sb(BUFFER_SIZE); \ + TKafkaWritable writable(sb); \ + TKafkaReadable readable(sb.GetFrontBuffer()); \ + \ + Y_UNUSED(readable); \ + Y_UNUSED(result); \ + \ + NKafka::NPrivate::TWriteCollector collector; + +struct Meta_TKafkaInt8 { + using Type = TKafkaInt8; + using TypeDesc = NKafka::NPrivate::TKafkaIntDesc; + + static constexpr const char* Name = "value"; + static constexpr const char* About = "The test field."; + static constexpr const TKafkaInt32 Tag = 31; + static const Type Default; + + static constexpr TKafkaVersions PresentVersions{3, 97}; + static constexpr TKafkaVersions TaggedVersions{11, 17}; + static constexpr TKafkaVersions NullableVersions{5, 19}; + static constexpr TKafkaVersions FlexibleVersions{7, Max()}; +}; + +const Meta_TKafkaInt8::Type Meta_TKafkaInt8::Default = 7; + +static_assert(Meta_TKafkaInt8::Tag == 31); +static_assert(Meta_TKafkaInt8::About[0] != '\0'); +static_assert(Meta_TKafkaInt8::NullableVersions.Min == 5); +static_assert(Meta_TKafkaInt8::FlexibleVersions.Min == 7); + +struct Meta_TKafkaString { + using Type = TKafkaString; + using TypeDesc = NKafka::NPrivate::TKafkaStringDesc; + + static constexpr const char* Name = "value"; + static constexpr const char* About = "The test field."; + static constexpr const TKafkaInt32 Tag = 31; + static const Type Default; + + static constexpr TKafkaVersions PresentVersions{3, 97}; + static constexpr TKafkaVersions TaggedVersions{11, 17}; + static constexpr TKafkaVersions NullableVersions{5, 19}; + static constexpr TKafkaVersions FlexibleVersions{7, Max()}; +}; + +const Meta_TKafkaString::Type Meta_TKafkaString::Default = "default_value"; + +static_assert(Meta_TKafkaString::Tag == 31); +static_assert(Meta_TKafkaString::About[0] != '\0'); +static_assert(Meta_TKafkaString::NullableVersions.Min == 5); +static_assert(Meta_TKafkaString::FlexibleVersions.Min == 7); + +struct Meta_TKafkaArray { + using Type = std::vector; + using TypeDesc = NKafka::NPrivate::TKafkaArrayDesc; + using ItemType = TKafkaString; + using ItemTypeDesc = NKafka::NPrivate::TKafkaStringDesc; + + static constexpr const char* Name = "value"; + static constexpr const char* About = "The test field."; + static constexpr const TKafkaInt32 Tag = 31; + + static constexpr TKafkaVersions PresentVersions{3, 97}; + static constexpr TKafkaVersions TaggedVersions{11, 17}; + static constexpr TKafkaVersions NullableVersions{5, 19}; + static constexpr TKafkaVersions FlexibleVersions{7, Max()}; +}; + +static_assert(Meta_TKafkaArray::Tag == 31); +static_assert(Meta_TKafkaArray::About[0] != '\0'); +static_assert(Meta_TKafkaArray::NullableVersions.Min == 5); +static_assert(Meta_TKafkaArray::FlexibleVersions.Min == 7); + +struct Meta_TKafkaBytes { + using Type = TKafkaBytes; + using TypeDesc = NKafka::NPrivate::TKafkaBytesDesc; + + static constexpr const char* Name = "value"; + static constexpr const char* About = "The test field."; + static constexpr const TKafkaInt32 Tag = 31; + + static constexpr TKafkaVersions PresentVersions{3, 97}; + static constexpr TKafkaVersions TaggedVersions{11, 17}; + static constexpr TKafkaVersions NullableVersions{5, 19}; + static constexpr TKafkaVersions FlexibleVersions{7, Max()}; +}; + +static_assert(Meta_TKafkaBytes::Tag == 31); +static_assert(Meta_TKafkaBytes::About[0] != '\0'); +static_assert(Meta_TKafkaBytes::NullableVersions.Min == 5); +static_assert(Meta_TKafkaBytes::FlexibleVersions.Min == 7); + +struct Meta_TKafkaFloat64 { + using Type = TKafkaFloat64; + using TypeDesc = NKafka::NPrivate::TKafkaFloat64Desc; + + static constexpr const char* Name = "value"; + static constexpr const char* About = "The test field."; + static constexpr const TKafkaInt32 Tag = 31; + static const Type Default; + + static constexpr TKafkaVersions PresentVersions{3, 97}; + static constexpr TKafkaVersions TaggedVersions{11, 17}; + static constexpr TKafkaVersions NullableVersions{5, 19}; + static constexpr TKafkaVersions FlexibleVersions{7, Max()}; +}; + +const Meta_TKafkaFloat64::Type Meta_TKafkaFloat64::Default = 7.875; + +static_assert(Meta_TKafkaFloat64::Tag == 31); +static_assert(Meta_TKafkaFloat64::About[0] != '\0'); +static_assert(Meta_TKafkaFloat64::NullableVersions.Min == 5); +static_assert(Meta_TKafkaFloat64::FlexibleVersions.Min == 7); + +Y_UNIT_TEST_SUITE(KafkaMessagesInt) { + Y_UNIT_TEST(TKafkaInt8_NotPresentVersion) { + SIMPLE_HEAD(TKafkaInt8, 37); + + NKafka::NPrivate::Write(collector, writable, 0, value); + UNIT_ASSERT_EQUAL(sb.GetFrontBuffer().size(), (size_t)0); + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 0u); + + NKafka::NPrivate::Read(readable, 0, result); + UNIT_ASSERT_EQUAL(result, Meta_TKafkaInt8::Default); + } + + Y_UNIT_TEST(TKafkaInt8_PresentVersion_NotTaggedVersion) { + SIMPLE_HEAD(TKafkaInt8, 37); + + NKafka::NPrivate::Write(collector, writable, 3, value); + NKafka::NPrivate::Read(readable, 3, result); + + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 0u); + UNIT_ASSERT_EQUAL(result, value); + } + + Y_UNIT_TEST(TKafkaInt8_PresentVersion_TaggedVersion) { + SIMPLE_HEAD(TKafkaInt8, 37); + + NKafka::NPrivate::Write(collector, writable, 11, value); + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 1u); + + NKafka::NPrivate::WriteTag(writable, 11, value); + + ui32 tag = readable.readUnsignedVarint(); + UNIT_ASSERT_EQUAL(tag, Meta_TKafkaInt8::Tag); + + ui32 size = readable.readUnsignedVarint(); + UNIT_ASSERT_EQUAL(size, sizeof(TKafkaInt8)); + + NKafka::NPrivate::ReadTag(readable, 11, result); + UNIT_ASSERT_EQUAL(result, value); + } + + Y_UNIT_TEST(TKafkaInt8_PresentVersion_TaggedVersion_Default) { + SIMPLE_HEAD(TKafkaInt8, Meta_TKafkaInt8::Default); + + NKafka::NPrivate::Write(collector, writable, 11, value); + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 0u); + } + + Y_UNIT_TEST(TKafkaString_IsDefault) { + TKafkaString value; + UNIT_ASSERT(!NKafka::NPrivate::IsDefaultValue(value)); + + value = "random_string"; + UNIT_ASSERT(!NKafka::NPrivate::IsDefaultValue(value)); + + value = "default_value"; + UNIT_ASSERT(NKafka::NPrivate::IsDefaultValue(value)); + } + + Y_UNIT_TEST(TKafkaString_PresentVersion_NotTaggedVersion) { + SIMPLE_HEAD(TKafkaString, { "some value" }); + + NKafka::NPrivate::Write(collector, writable, 3, value); + NKafka::NPrivate::Read(readable, 3, result); + + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 0u); + UNIT_ASSERT_EQUAL(result, value); + } + + Y_UNIT_TEST(TKafkaString_PresentVersion_TaggedVersion) { + SIMPLE_HEAD(TKafkaString, { "some value" }); + + NKafka::NPrivate::Write(collector, writable, 11, value); + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 1u); + + NKafka::NPrivate::WriteTag(writable, 11, value); + + ui32 tag = readable.readUnsignedVarint(); + UNIT_ASSERT_EQUAL(tag, Meta_TKafkaString::Tag); + + ui32 size = readable.readUnsignedVarint(); + UNIT_ASSERT_EQUAL(size, value->size() + NKafka::NPrivate::SizeOfUnsignedVarint(value->size() + 1)); + + NKafka::NPrivate::ReadTag(readable, 11, result); + UNIT_ASSERT_EQUAL(result, value); + } + + Y_UNIT_TEST(TKafkaString_PresentVersion_TaggedVersion_Default) { + SIMPLE_HEAD(TKafkaString, Meta_TKafkaString::Default); + + NKafka::NPrivate::Write(collector, writable, 11, value); + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 0u); + } + + Y_UNIT_TEST(TKafkaArray_IsDefault) { + Meta_TKafkaArray::Type value; + UNIT_ASSERT(NKafka::NPrivate::IsDefaultValue(value)); + + value.push_back("random_string"); + UNIT_ASSERT(!NKafka::NPrivate::IsDefaultValue(value)); + } + + Y_UNIT_TEST(TKafkaArray_PresentVersion_NotTaggedVersion) { + SIMPLE_HEAD(TKafkaArray, { "some value" }); + + NKafka::NPrivate::Write(collector, writable, 3, value); + NKafka::NPrivate::Read(readable, 3, result); + + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 0u); + UNIT_ASSERT_EQUAL(result, value); + } + + Y_UNIT_TEST(TKafkaArray_PresentVersion_TaggedVersion) { + TString v = "some value"; + SIMPLE_HEAD(TKafkaArray, { v }); + + NKafka::NPrivate::Write(collector, writable, 11, value); + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 1u); + + NKafka::NPrivate::WriteTag(writable, 11, value); + + ui32 tag = readable.readUnsignedVarint(); + UNIT_ASSERT_EQUAL(tag, Meta_TKafkaArray::Tag); + + ui32 size = readable.readUnsignedVarint(); + UNIT_ASSERT_EQUAL(size, v.length() + + NKafka::NPrivate::SizeOfUnsignedVarint(value.size()) + + NKafka::NPrivate::SizeOfUnsignedVarint(v.length() + 1)); + + NKafka::NPrivate::ReadTag(readable, 11, result); + UNIT_ASSERT_EQUAL(result, value); + } + + Y_UNIT_TEST(TKafkaArray_PresentVersion_TaggedVersion_Default) { + SIMPLE_HEAD(TKafkaArray, {}); + + NKafka::NPrivate::Write(collector, writable, 11, value); + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 0u); + } + + Y_UNIT_TEST(TKafkaBytes_IsDefault) { + Meta_TKafkaBytes::Type value; + UNIT_ASSERT(NKafka::NPrivate::IsDefaultValue(value)); + + char v[] = "value"; + value = TArrayRef(v); + UNIT_ASSERT(!NKafka::NPrivate::IsDefaultValue(value)); + } + + Y_UNIT_TEST(TKafkaBytes_PresentVersion_NotTaggedVersion) { + char v[] = "0123456789"; + SIMPLE_HEAD(TKafkaBytes, TArrayRef(v)); + + NKafka::NPrivate::Write(collector, writable, 3, value); + NKafka::NPrivate::Read(readable, 3, result); + + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 0u); + UNIT_ASSERT_EQUAL(*result, *value); + } + + Y_UNIT_TEST(TKafkaBytes_PresentVersion_TaggedVersion) { + char v[] = "0123456789"; + SIMPLE_HEAD(TKafkaBytes, TArrayRef(v)); + + NKafka::NPrivate::Write(collector, writable, 11, value); + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 1u); + + NKafka::NPrivate::WriteTag(writable, 11, value); + + ui32 tag = readable.readUnsignedVarint(); + UNIT_ASSERT_EQUAL(tag, Meta_TKafkaBytes::Tag); + + ui32 size = readable.readUnsignedVarint(); + UNIT_ASSERT_EQUAL(size, value->size() + + NKafka::NPrivate::SizeOfUnsignedVarint(value->size() + 1)); + + NKafka::NPrivate::ReadTag(readable, 11, result); + UNIT_ASSERT_EQUAL(*result, *value); + } + + Y_UNIT_TEST(TKafkaBytes_PresentVersion_TaggedVersion_Default) { + SIMPLE_HEAD(TKafkaBytes, std::nullopt); + + NKafka::NPrivate::Write(collector, writable, 11, value); + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 0u); + } + + Y_UNIT_TEST(TKafkaFloat64_PresentVersion_NotTaggedVersion) { + ui8 reference[] = {0x40, 0x09, 0x21, 0xCA, 0xC0, 0x83, 0x12, 0x6F}; + + SIMPLE_HEAD(TKafkaFloat64, 3.1415); + + NKafka::NPrivate::Write(collector, writable, 3, value); + NKafka::NPrivate::Read(readable, 3, result); + + UNIT_ASSERT_EQUAL(collector.NumTaggedFields, 0u); + UNIT_ASSERT_EQUAL(result, value); + + UNIT_ASSERT_EQUAL(sb.GetFrontBuffer().size(), sizeof(reference)); + for (size_t i = 0; i < sizeof(reference); ++i) { + UNIT_ASSERT_EQUAL(*(sb.GetFrontBuffer().data() + i), (char)reference[i]); + } + } +} + +} // namespace +} // namespace NKafka diff --git a/src/library/kafka/ut/kafka_records_ut.cpp b/src/library/kafka/ut/kafka_records_ut.cpp new file mode 100644 index 00000000000..355b10184cd --- /dev/null +++ b/src/library/kafka/ut/kafka_records_ut.cpp @@ -0,0 +1,529 @@ +#include +#include +#include + +#include + +namespace NKafka { +namespace { + +static constexpr size_t BUFFER_SIZE = 1 << 16; + +template +void CheckUnsignedVarint(const std::vector& values) { + for (T v : values) { + Cerr << ">>>>> Check value=" << v << Endl << Flush; + TKafkaWriteBuffer sb(BUFFER_SIZE); + TKafkaWritable writable(sb); + TKafkaReadable readable(sb.GetFrontBuffer()); + + writable.writeUnsignedVarint(v); + + UNIT_ASSERT_EQUAL_C(sb.GetFrontBuffer().size(), NKafka::NPrivate::SizeOfUnsignedVarint(v), + TStringBuilder() << "Size mismatch " << sb.GetFrontBuffer().size() << " != " << NKafka::NPrivate::SizeOfUnsignedVarint(v)); + + T r = readable.readUnsignedVarint(); + UNIT_ASSERT_EQUAL_C(r, v, TStringBuilder() << r << " != " << v); + } +} + +template +void CheckVarint(const std::vector& values) { + for (T v : values) { + Cerr << ">>>>> Check value=" << v << Endl << Flush; + TKafkaWriteBuffer sb(BUFFER_SIZE); + TKafkaWritable writable(sb); + TKafkaReadable readable(sb.GetFrontBuffer()); + + writable.writeVarint(v); + + UNIT_ASSERT_EQUAL_C(sb.GetFrontBuffer().size(), NKafka::NPrivate::SizeOfVarint(v), + TStringBuilder() << "Size mismatch " << sb.GetFrontBuffer().size() << " != " << NKafka::NPrivate::SizeOfVarint(v)); + + T r = readable.readVarint(); + + UNIT_ASSERT_EQUAL_C(r, v, TStringBuilder() << r << " != " << v); + } +} + +template +void CheckVarintWrongBytes(std::vector bytes) { + TKafkaWriteBuffer sb(BUFFER_SIZE); + TKafkaWritable writable(sb); + TKafkaReadable readable(sb.GetFrontBuffer()); + + writable.write(reinterpret_cast(bytes.data()), bytes.size()); + + try { + readable.readUnsignedVarint(); + UNIT_FAIL("Must be exception"); + } catch (const yexception& e) { + UNIT_ASSERT_STRING_CONTAINS(e.what(), "illegal varint length"); + } +} + +TKafkaRecord MakeRecord(i64 timestampDelta, i64 offsetDelta, TStringBuf key, TStringBuf value) { + TKafkaRecord record; + record.TimestampDelta = timestampDelta; + record.OffsetDelta = offsetDelta; + record.Key = key; + record.Value = value; + record.Length = record.Size(2) - NKafka::NPrivate::SizeOfVarint(0); + return record; +} + +TKafkaRecordBatch MakeRecordBatch(ECompressionType compressionType) { + TKafkaRecordBatch batch; + batch.BaseOffset = 42; + batch.Magic = 2; + batch.Attributes = static_cast(compressionType); + batch.BaseSequence = 17; + batch.LastOffsetDelta = 1; + batch.BaseTimestamp = 1000; + batch.MaxTimestamp = 1010; + batch.Records.push_back(MakeRecord(7, 0, "key-0", "value-0")); + batch.Records.push_back(MakeRecord(10, 1, "key-1", "value-1")); + batch.BatchLength = batch.Size(2) - sizeof(TKafkaRecordBatch::BaseOffsetMeta::Type) - sizeof(TKafkaRecordBatch::BatchLengthMeta::Type); + return batch; +} + +void AssertRecordBatchRoundTrip(ECompressionType compressionType) { + const TKafkaRecordBatch batch = MakeRecordBatch(compressionType); + + const TString serialized = WriteKafkaRecordBatch(batch); + TKafkaRecordBatch parsed = ReadKafkaRecordBatch(serialized); + + UNIT_ASSERT_VALUES_EQUAL(parsed.BaseOffset, batch.BaseOffset); + UNIT_ASSERT_VALUES_EQUAL(parsed.Magic, batch.Magic); + UNIT_ASSERT_VALUES_EQUAL(static_cast(parsed.CompressionType()), static_cast(compressionType)); + UNIT_ASSERT_VALUES_EQUAL(parsed.BatchLength, batch.BatchLength); + UNIT_ASSERT_VALUES_EQUAL(parsed.Records.size(), batch.Records.size()); + for (size_t i = 0; i < batch.Records.size(); ++i) { + UNIT_ASSERT_VALUES_EQUAL(parsed.Records[i].TimestampDelta, batch.Records[i].TimestampDelta); + UNIT_ASSERT_VALUES_EQUAL(parsed.Records[i].OffsetDelta, batch.Records[i].OffsetDelta); + UNIT_ASSERT(KafkaBytesEqual(parsed.Records[i].Key, batch.Records[i].Key)); + UNIT_ASSERT(KafkaBytesEqual(parsed.Records[i].Value, batch.Records[i].Value)); + } +} + +void AssertUnsupportedCompressionType(ECompressionType compressionType) { + TKafkaRecordBatch batch = MakeRecordBatch(ECompressionType::NONE); + batch.Attributes = static_cast(compressionType); + + try { + Y_UNUSED(batch.Size(2)); + UNIT_FAIL("Must be exception"); + } catch (const yexception& e) { + UNIT_ASSERT_STRING_CONTAINS(e.what(), "unsupported Kafka record batch compression type"); + } +} + +TString Bytes(std::initializer_list bytes) { + TString result; + result.reserve(bytes.size()); + for (const ui8 byte : bytes) { + result.push_back(static_cast(byte)); + } + return result; +} + +// Generated with Apache Kafka Java code: +// +// for (CompressionType type : new CompressionType[]{CompressionType.NONE, CompressionType.GZIP, CompressionType.ZSTD}) { +// MemoryRecordsBuilder builder = MemoryRecords.builder( +// ByteBuffer.allocate(1024), RecordBatch.MAGIC_VALUE_V2, Compression.of(type).build(), +// TimestampType.CREATE_TIME, 42L, RecordBatch.NO_TIMESTAMP); +// ProducerBatch batch = new ProducerBatch(new TopicPartition("topic", 0), builder, 1000L); +// batch.tryAppend(1007L, "key-0".getBytes(), "value-0".getBytes(), +// new Header[]{new RecordHeader("header-0", "hvalue-0".getBytes())}, null, 1007L); +// batch.tryAppend(1010L, "key-1".getBytes(), "value-1".getBytes(), +// new Header[]{new RecordHeader("header-1", "hvalue-1".getBytes())}, null, 1010L); +// ByteBuffer buffer = batch.records().buffer(); +// byte[] bytes = new byte[buffer.remaining()]; +// buffer.get(bytes); +// System.out.println(type + " " + toCppBytes(bytes)); +// } +TString KafkaProducerBatchBytes(ECompressionType compressionType) { + switch (compressionType) { + case ECompressionType::NONE: + return Bytes({ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2A, 0x00, 0x00, 0x00, 0x7B, 0xFF, 0xFF, 0xFF, 0xFF, + 0x02, 0xF6, 0x09, 0xC5, 0x43, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x03, 0xEF, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xF2, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, + 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x02, 0x48, 0x00, 0x00, + 0x00, 0x0A, 0x6B, 0x65, 0x79, 0x2D, 0x30, 0x0E, 0x76, 0x61, 0x6C, 0x75, 0x65, 0x2D, 0x30, 0x02, + 0x10, 0x68, 0x65, 0x61, 0x64, 0x65, 0x72, 0x2D, 0x30, 0x10, 0x68, 0x76, 0x61, 0x6C, 0x75, 0x65, + 0x2D, 0x30, 0x48, 0x00, 0x06, 0x02, 0x0A, 0x6B, 0x65, 0x79, 0x2D, 0x31, 0x0E, 0x76, 0x61, 0x6C, + 0x75, 0x65, 0x2D, 0x31, 0x02, 0x10, 0x68, 0x65, 0x61, 0x64, 0x65, 0x72, 0x2D, 0x31, 0x10, 0x68, + 0x76, 0x61, 0x6C, 0x75, 0x65, 0x2D, 0x31, + }); + case ECompressionType::GZIP: + return Bytes({ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2A, 0x00, 0x00, 0x00, 0x75, 0xFF, 0xFF, 0xFF, 0xFF, + 0x02, 0x29, 0x52, 0xC6, 0x73, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x03, 0xEF, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xF2, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, + 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x02, 0x1F, 0x8B, 0x08, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xF3, 0x60, 0x60, 0x60, 0xE0, 0xCA, 0x4E, 0xAD, 0xD4, + 0x35, 0xE0, 0x2B, 0x4B, 0xCC, 0x29, 0x4D, 0xD5, 0x35, 0x60, 0x12, 0xC8, 0x48, 0x4D, 0x4C, 0x49, + 0x2D, 0xD2, 0x35, 0x10, 0xC8, 0x80, 0x0A, 0x79, 0x30, 0xB0, 0x31, 0x81, 0x15, 0x19, 0x42, 0x15, + 0x19, 0xC2, 0x15, 0x19, 0xC2, 0x14, 0x19, 0x02, 0x00, 0xB1, 0x38, 0x5C, 0x0F, 0x4A, 0x00, 0x00, + 0x00, + }); + case ECompressionType::ZSTD: + return Bytes({ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2A, 0x00, 0x00, 0x00, 0x72, 0xFF, 0xFF, 0xFF, 0xFF, + 0x02, 0xD6, 0xD3, 0xEA, 0xF9, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x03, 0xEF, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xF2, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, + 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x02, 0x28, 0xB5, 0x2F, + 0xFD, 0x00, 0x58, 0xAC, 0x01, 0x00, 0x64, 0x02, 0x48, 0x00, 0x00, 0x00, 0x0A, 0x6B, 0x65, 0x79, + 0x2D, 0x30, 0x0E, 0x76, 0x61, 0x6C, 0x75, 0x65, 0x2D, 0x30, 0x02, 0x10, 0x68, 0x65, 0x61, 0x64, + 0x65, 0x72, 0x2D, 0x30, 0x10, 0x68, 0x48, 0x00, 0x06, 0x02, 0x31, 0x31, 0x31, 0x31, 0x05, 0x00, + 0x80, 0x08, 0x28, 0x07, 0x92, 0x21, 0xD5, 0xB2, 0xE6, 0x4E, 0x27, 0x01, 0x00, 0x00, + }); + default: + ythrow yexception() << "unexpected golden Kafka record batch compression type: " << static_cast(compressionType); + } +} + +// Generated with Apache Kafka Java code: +// +// for (byte magic : new byte[]{RecordBatch.MAGIC_VALUE_V0, RecordBatch.MAGIC_VALUE_V1}) { +// for (CompressionType type : new CompressionType[]{CompressionType.NONE, CompressionType.GZIP}) { +// MemoryRecords records = MemoryRecords.withRecords( +// magic, 42L, Compression.of(type).build(), TimestampType.CREATE_TIME, +// new SimpleRecord(1007L, "key-0".getBytes(), "value-0".getBytes()), +// new SimpleRecord(1010L, "key-1".getBytes(), "value-1".getBytes())); +// ByteBuffer buffer = records.buffer(); +// byte[] bytes = new byte[buffer.remaining()]; +// buffer.get(bytes); +// System.out.println("magic" + magic + "-" + type + " " + toCppBytes(bytes)); +// } +// } +TString KafkaLegacyProducerBatchBytes(TKafkaVersion magic, ECompressionType compressionType) { + switch (magic) { + case 0: + switch (compressionType) { + case ECompressionType::NONE: + return Bytes({ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2A, 0x00, 0x00, 0x00, 0x1A, 0xE4, 0xFF, 0xC8, 0x29, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x05, 0x6B, 0x65, 0x79, 0x2D, 0x30, 0x00, 0x00, 0x00, 0x07, 0x76, + 0x61, 0x6C, 0x75, 0x65, 0x2D, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2B, 0x00, 0x00, + 0x00, 0x1A, 0x08, 0x5D, 0xB4, 0xD0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x05, 0x6B, 0x65, 0x79, 0x2D, + 0x31, 0x00, 0x00, 0x00, 0x07, 0x76, 0x61, 0x6C, 0x75, 0x65, 0x2D, 0x31, + }); + case ECompressionType::GZIP: + return Bytes({ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2B, 0x00, 0x00, 0x00, 0x4D, 0xB0, 0xC5, 0x48, 0xB4, + 0x00, 0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, 0x00, 0x3F, 0x1F, 0x8B, 0x08, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0xFF, 0x63, 0x60, 0x00, 0x03, 0x2D, 0x20, 0x96, 0x7A, 0xF2, 0xFF, 0x84, 0x26, + 0x98, 0xC7, 0x9A, 0x9D, 0x5A, 0xA9, 0x6B, 0x00, 0x64, 0xB0, 0x97, 0x25, 0xE6, 0x94, 0xA6, 0x82, + 0x99, 0x20, 0xA0, 0x0D, 0x52, 0xC5, 0x11, 0xBB, 0xE5, 0x02, 0x42, 0x95, 0x21, 0x42, 0x95, 0x21, + 0x00, 0xB1, 0xFA, 0xE3, 0x39, 0x4C, 0x00, 0x00, 0x00, + }); + default: + break; + } + break; + case 1: + switch (compressionType) { + case ECompressionType::NONE: + return Bytes({ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2A, 0x00, 0x00, 0x00, 0x22, 0x54, 0x45, 0x73, 0x90, + 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xEF, 0x00, 0x00, 0x00, 0x05, 0x6B, 0x65, + 0x79, 0x2D, 0x30, 0x00, 0x00, 0x00, 0x07, 0x76, 0x61, 0x6C, 0x75, 0x65, 0x2D, 0x30, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x2B, 0x00, 0x00, 0x00, 0x22, 0x5B, 0xD2, 0x93, 0xD6, 0x01, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xF2, 0x00, 0x00, 0x00, 0x05, 0x6B, 0x65, 0x79, 0x2D, + 0x31, 0x00, 0x00, 0x00, 0x07, 0x76, 0x61, 0x6C, 0x75, 0x65, 0x2D, 0x31, + }); + case ECompressionType::GZIP: + return Bytes({ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2B, 0x00, 0x00, 0x00, 0x5A, 0xB0, 0x86, 0xC9, 0x13, + 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03, 0xF2, 0xFF, 0xFF, 0xFF, 0xFF, 0x00, 0x00, + 0x00, 0x44, 0x1F, 0x8B, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0x63, 0x60, 0x80, 0x03, + 0xA5, 0x10, 0xD7, 0xE2, 0x09, 0x8C, 0x50, 0x0E, 0xF3, 0x7B, 0x20, 0xC1, 0x9A, 0x9D, 0x5A, 0xA9, + 0x6B, 0x00, 0x64, 0xB0, 0x97, 0x25, 0xE6, 0x94, 0xA6, 0x82, 0x99, 0x20, 0x00, 0x52, 0xA4, 0x14, + 0x7D, 0x69, 0xF2, 0x35, 0xB8, 0xEA, 0x4F, 0x30, 0xD5, 0x86, 0x08, 0xD5, 0x86, 0x00, 0xE6, 0x0F, + 0xD7, 0x68, 0x5C, 0x00, 0x00, 0x00, + }); + default: + break; + } + break; + default: + break; + } + + ythrow yexception() << "unexpected golden Kafka legacy record batch magic/compression: " + << magic << "/" << static_cast(compressionType); +} + +void AssertKafkaRecord( + const TKafkaRecord& record, + i64 timestampDelta, + i64 offsetDelta, + TStringBuf key, + TStringBuf value, + TStringBuf headerKey, + TStringBuf headerValue) +{ + UNIT_ASSERT_VALUES_EQUAL(record.TimestampDelta, timestampDelta); + UNIT_ASSERT_VALUES_EQUAL(record.OffsetDelta, offsetDelta); + UNIT_ASSERT(record.Key); + UNIT_ASSERT(record.Value); + UNIT_ASSERT(KafkaBytesEqual(record.Key, key)); + UNIT_ASSERT(KafkaBytesEqual(record.Value, value)); + UNIT_ASSERT_VALUES_EQUAL(record.Headers.size(), 1); + UNIT_ASSERT(record.Headers[0].Key); + UNIT_ASSERT(record.Headers[0].Value); + UNIT_ASSERT(KafkaBytesEqual(record.Headers[0].Key, headerKey)); + UNIT_ASSERT(KafkaBytesEqual(record.Headers[0].Value, headerValue)); +} + +void AssertKafkaLegacyRecord( + const TKafkaRecord& record, + i64 timestampDelta, + i64 offsetDelta, + TStringBuf key, + TStringBuf value) +{ + UNIT_ASSERT_VALUES_EQUAL(record.TimestampDelta, timestampDelta); + UNIT_ASSERT_VALUES_EQUAL(record.OffsetDelta, offsetDelta); + UNIT_ASSERT(record.Key); + UNIT_ASSERT(record.Value); + UNIT_ASSERT(KafkaBytesEqual(record.Key, key)); + UNIT_ASSERT(KafkaBytesEqual(record.Value, value)); + UNIT_ASSERT(record.Headers.empty()); +} + +void AssertKafkaProducerBatchDeserialized(ECompressionType compressionType) { + TKafkaRecordBatch parsed = ReadKafkaRecordBatch(KafkaProducerBatchBytes(compressionType)); + + UNIT_ASSERT_VALUES_EQUAL(parsed.BaseOffset, 42); + UNIT_ASSERT_VALUES_EQUAL(parsed.PartitionLeaderEpoch, -1); + UNIT_ASSERT_VALUES_EQUAL(parsed.Magic, 2); + UNIT_ASSERT_VALUES_EQUAL(static_cast(parsed.CompressionType()), static_cast(compressionType)); + UNIT_ASSERT_VALUES_EQUAL(parsed.LastOffsetDelta, 1); + UNIT_ASSERT_VALUES_EQUAL(parsed.BaseTimestamp, 1007); + UNIT_ASSERT_VALUES_EQUAL(parsed.MaxTimestamp, 1010); + UNIT_ASSERT_VALUES_EQUAL(parsed.ProducerId, -1); + UNIT_ASSERT_VALUES_EQUAL(parsed.ProducerEpoch, -1); + UNIT_ASSERT_VALUES_EQUAL(parsed.BaseSequence, -1); + UNIT_ASSERT_VALUES_EQUAL(parsed.Records.size(), 2); + + AssertKafkaRecord(parsed.Records[0], 0, 0, "key-0", "value-0", "header-0", "hvalue-0"); + AssertKafkaRecord(parsed.Records[1], 3, 1, "key-1", "value-1", "header-1", "hvalue-1"); +} + +void AssertKafkaProducerBatchSerialized(ECompressionType compressionType) { + const TString serialized = KafkaProducerBatchBytes(compressionType); + TKafkaRecordBatch parsed = ReadKafkaRecordBatch( + serialized, + 2); + + if (compressionType == ECompressionType::NONE) { + UNIT_ASSERT_VALUES_EQUAL(WriteKafkaRecordBatch(parsed), serialized); + } else { + const TString roundTrip = WriteKafkaRecordBatch(parsed); + TKafkaRecordBatch reparsed = ReadKafkaRecordBatch(roundTrip); + UNIT_ASSERT_VALUES_EQUAL(reparsed.Records.size(), parsed.Records.size()); + } +} + +TKafkaRecordBatch ReadKafkaLegacyProducerBatch(TStringBuf data, TKafkaVersion magic) { + TBuffer buffer(data.data(), data.size()); + TKafkaReadable readable(buffer); + readable.SetAllowCompressed(true); + + TKafkaRecordBatch parsed; + NPrivate::ReadLegacyRecordBatch(readable, magic, data.size(), parsed); + UNIT_ASSERT_VALUES_EQUAL(readable.left(), 0); + return parsed; +} + +// TKafkaRecordBatchV0 keeps TKafkaBytes views into `buffer`, so the caller must +// keep `buffer` alive for as long as the returned records are used. +std::vector ReadKafkaLegacyRecordBatchWrappers(const TBuffer& buffer, TKafkaVersion magic) { + TKafkaReadable readable(buffer); + std::vector records; + while (readable.left() > 0) { + UNIT_ASSERT_VALUES_EQUAL(readable.take(16), magic); + auto& record = records.emplace_back(); + record.Read(readable, magic); + } + return records; +} + +TString WriteKafkaLegacyRecordBatchWrappers(const std::vector& records, TKafkaVersion magic) { + size_t size = 0; + for (const auto& record : records) { + size += record.Size(magic); + } + + TKafkaWriteBuffer buffer(size); + TKafkaWritable writable(buffer); + for (const auto& record : records) { + record.Write(writable, magic); + } + return buffer.AsString(); +} + +void AssertKafkaLegacyProducerBatchDeserialized( + TKafkaVersion magic, + ECompressionType compressionType, + i64 expectedFirstTimestamp, + i64 expectedSecondTimestamp) +{ + TKafkaRecordBatch parsed = ReadKafkaLegacyProducerBatch( + KafkaLegacyProducerBatchBytes(magic, compressionType), + magic); + + UNIT_ASSERT_VALUES_EQUAL(static_cast(parsed.CompressionType()), static_cast(compressionType)); + UNIT_ASSERT_VALUES_EQUAL(parsed.Records.size(), 2); + AssertKafkaLegacyRecord(parsed.Records[0], expectedFirstTimestamp, 42, "key-0", "value-0"); + AssertKafkaLegacyRecord(parsed.Records[1], expectedSecondTimestamp, 43, "key-1", "value-1"); +} + +void AssertKafkaLegacyProducerBatchSerialized(TKafkaVersion magic, ECompressionType compressionType) { + const TString serialized = KafkaLegacyProducerBatchBytes(magic, compressionType); + TBuffer buffer(serialized.data(), serialized.size()); + const std::vector records = ReadKafkaLegacyRecordBatchWrappers(buffer, magic); + UNIT_ASSERT_VALUES_EQUAL(WriteKafkaLegacyRecordBatchWrappers(records, magic), serialized); +} + +void AssertReadBatchHeaderForProducerBatch(ECompressionType compressionType) { + const TString bytes = KafkaProducerBatchBytes(compressionType); + const auto header = ReadKafkaBatchHeader(bytes); + UNIT_ASSERT(header); + UNIT_ASSERT_VALUES_EQUAL(header->RecordsCount, ReadKafkaRecordBatch(bytes).Records.size()); + UNIT_ASSERT_VALUES_EQUAL(header->BaseSequence, -1); +} + +void AssertReadBatchHeaderForLegacyProducerBatch(TKafkaVersion magic, ECompressionType compressionType) { + const TString bytes = KafkaLegacyProducerBatchBytes(magic, compressionType); + const auto header = ReadKafkaBatchHeader(bytes); + UNIT_ASSERT(header); + UNIT_ASSERT_VALUES_EQUAL(header->Magic, magic); + UNIT_ASSERT_VALUES_EQUAL(header->RecordsCount, 2); + UNIT_ASSERT_VALUES_EQUAL(ReadKafkaLegacyProducerBatch(bytes, magic).Records.size(), 2u); +} + +Y_UNIT_TEST_SUITE(KafkaRecords) { + Y_UNIT_TEST(UnsignedVarint32) { + CheckUnsignedVarint({0, 1, 127, 128, 32191, Max(), Max()}); + } + + Y_UNIT_TEST(UnsignedVarint64) { + CheckUnsignedVarint({0, 1, 127, 128, 32191, Max(), static_cast(Max()) + 1, Max(), Max()}); + } + + Y_UNIT_TEST(Varint32) { + CheckVarint({Min(), -167966, -1, 0, 1, 127, 128, 32191, Max()}); + } + + Y_UNIT_TEST(Varint64) { + CheckVarint({Min(), Min(), -167966, -1, 0, 1, 127, 128, 32191, static_cast(Max()) + 1, Max()}); + } + + Y_UNIT_TEST(UnsignedVarint32Wrong) { + CheckVarintWrongBytes({0xFF, 0xFF, 0xFF, 0xFF, 0xFF}); + } + + Y_UNIT_TEST(UnsignedVarint64Wrong) { + CheckVarintWrongBytes({0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF}); + } + + Y_UNIT_TEST(UnsignedVarint32Deserialize) { + std::vector bytes = {0x81, 0x83, 0x05}; + + TKafkaWriteBuffer sb(BUFFER_SIZE); + TKafkaWritable writable(sb); + TKafkaReadable readable(sb.GetFrontBuffer()); + + writable.write(reinterpret_cast(bytes.data()), bytes.size()); + + ui32 result = readable.readUnsignedVarint(); + UNIT_ASSERT_EQUAL(result, 1 + (3 << 7) + (5 << 14)); + } + + Y_UNIT_TEST(RecordBatchRoundTrip) { + AssertRecordBatchRoundTrip(ECompressionType::NONE); + } + + Y_UNIT_TEST(RecordBatchGzipRoundTrip) { + AssertRecordBatchRoundTrip(ECompressionType::GZIP); + } + + Y_UNIT_TEST(RecordBatchZstdRoundTrip) { + AssertRecordBatchRoundTrip(ECompressionType::ZSTD); + } + + Y_UNIT_TEST(ReadKafkaBatchHeader) { + UNIT_ASSERT(!ReadKafkaBatchHeader(TStringBuf())); + + const TKafkaRecordBatch batch = MakeRecordBatch(ECompressionType::NONE); + const TString batchBytes = WriteKafkaRecordBatch(batch); + const auto header = ReadKafkaBatchHeader(batchBytes); + UNIT_ASSERT(header); + UNIT_ASSERT_VALUES_EQUAL(header->RecordsCount, batch.Records.size()); + UNIT_ASSERT_VALUES_EQUAL(header->BaseSequence, batch.BaseSequence); + + const TKafkaRecordBatch gzipBatch = MakeRecordBatch(ECompressionType::GZIP); + const TString gzipBatchBytes = WriteKafkaRecordBatch(gzipBatch); + const auto gzipHeader = ReadKafkaBatchHeader(gzipBatchBytes); + UNIT_ASSERT(gzipHeader); + UNIT_ASSERT_VALUES_EQUAL(gzipHeader->RecordsCount, gzipBatch.Records.size()); + UNIT_ASSERT_VALUES_EQUAL(gzipHeader->BaseSequence, gzipBatch.BaseSequence); + + const TKafkaRecordBatch zstdBatch = MakeRecordBatch(ECompressionType::ZSTD); + const TString zstdBatchBytes = WriteKafkaRecordBatch(zstdBatch); + const auto zstdHeader = ReadKafkaBatchHeader(zstdBatchBytes); + UNIT_ASSERT(zstdHeader); + UNIT_ASSERT_VALUES_EQUAL(zstdHeader->RecordsCount, zstdBatch.Records.size()); + UNIT_ASSERT_VALUES_EQUAL(zstdHeader->BaseSequence, zstdBatch.BaseSequence); + } + + Y_UNIT_TEST(ReadKafkaBatchHeaderGoldenBytes) { + AssertReadBatchHeaderForProducerBatch(ECompressionType::NONE); + AssertReadBatchHeaderForProducerBatch(ECompressionType::GZIP); + AssertReadBatchHeaderForProducerBatch(ECompressionType::ZSTD); + + AssertReadBatchHeaderForLegacyProducerBatch(0, ECompressionType::NONE); + AssertReadBatchHeaderForLegacyProducerBatch(0, ECompressionType::GZIP); + AssertReadBatchHeaderForLegacyProducerBatch(1, ECompressionType::NONE); + AssertReadBatchHeaderForLegacyProducerBatch(1, ECompressionType::GZIP); + } + + Y_UNIT_TEST(KafkaProducerRecordBatchDeserialize) { + AssertKafkaProducerBatchDeserialized(ECompressionType::NONE); + AssertKafkaProducerBatchDeserialized(ECompressionType::GZIP); + AssertKafkaProducerBatchDeserialized(ECompressionType::ZSTD); + } + + Y_UNIT_TEST(KafkaProducerRecordBatchSerialize) { + AssertKafkaProducerBatchSerialized(ECompressionType::NONE); + AssertKafkaProducerBatchSerialized(ECompressionType::GZIP); + AssertKafkaProducerBatchSerialized(ECompressionType::ZSTD); + } + + Y_UNIT_TEST(KafkaLegacyProducerRecordBatchDeserialize) { + AssertKafkaLegacyProducerBatchDeserialized(0, ECompressionType::NONE, 0, 0); + AssertKafkaLegacyProducerBatchDeserialized(0, ECompressionType::GZIP, 0, 0); + AssertKafkaLegacyProducerBatchDeserialized(1, ECompressionType::NONE, 1007, 1010); + AssertKafkaLegacyProducerBatchDeserialized(1, ECompressionType::GZIP, 1010, 1010); + } + + Y_UNIT_TEST(KafkaLegacyProducerRecordBatchSerialize) { + AssertKafkaLegacyProducerBatchSerialized(0, ECompressionType::NONE); + AssertKafkaLegacyProducerBatchSerialized(0, ECompressionType::GZIP); + AssertKafkaLegacyProducerBatchSerialized(1, ECompressionType::NONE); + AssertKafkaLegacyProducerBatchSerialized(1, ECompressionType::GZIP); + } + + Y_UNIT_TEST(RecordBatchUnsupportedCompressionType) { + AssertUnsupportedCompressionType(ECompressionType::SNAPPY); + AssertUnsupportedCompressionType(ECompressionType::LZ4); + } +} + +} // namespace +} // namespace NKafka diff --git a/src/library/kafka/ut/kafka_ut_compare.h b/src/library/kafka/ut/kafka_ut_compare.h new file mode 100644 index 00000000000..4f3f1aecbf1 --- /dev/null +++ b/src/library/kafka/ut/kafka_ut_compare.h @@ -0,0 +1,21 @@ +#pragma once + +#include + +#include + +namespace NKafka { + +inline bool KafkaBytesEqual(const TKafkaBytes& bytes, TStringBuf expected) { + return bytes && TStringBuf(bytes->data(), bytes->size()) == expected; +} + +inline bool KafkaBytesEqual(const TKafkaBytes& lhs, const TKafkaBytes& rhs) { + if (!lhs || !rhs) { + return !lhs && !rhs; + } + return lhs->size() == rhs->size() + && memcmp(lhs->data(), rhs->data(), lhs->size()) == 0; +} + +} diff --git a/src/library/kafka/ut/ut_common.h b/src/library/kafka/ut/ut_common.h new file mode 100644 index 00000000000..844c0b357dd --- /dev/null +++ b/src/library/kafka/ut/ut_common.h @@ -0,0 +1,20 @@ +#pragma once + +#include + +#include + +namespace NKafka::NTest { + +inline void AssertKafkaBatchPayload(TStringBuf payload, size_t expectedRecordsCount, char expectedFill, size_t expectedDataSize) { + const auto batch = ReadKafkaRecordBatch(payload); + UNIT_ASSERT_VALUES_EQUAL(batch.Records.size(), expectedRecordsCount); + + for (const auto& record : batch.Records) { + UNIT_ASSERT_C(record.Value.has_value(), "Kafka batch record has no value"); + UNIT_ASSERT_VALUES_EQUAL(record.Value->size(), expectedDataSize); + UNIT_ASSERT_VALUES_EQUAL(TStringBuf(record.Value->data(), record.Value->size()), TString(expectedDataSize, expectedFill)); + } +} + +} // namespace NKafka::NTest diff --git a/src/library/operation_id/operation_id.cpp b/src/library/operation_id/operation_id.cpp index f2f0e4f6280..d4c78ff18b3 100644 --- a/src/library/operation_id/operation_id.cpp +++ b/src/library/operation_id/operation_id.cpp @@ -79,6 +79,9 @@ std::string ProtoToString(const Ydb::TOperationId& proto) { case Ydb::TOperationId::FULL_BACKUP: res << "ydb://fullbackup"; break; + case Ydb::TOperationId::ANALYZE: + res << "ydb://analyze"; + break; default: Y_ABORT_UNLESS(false, "unexpected kind"); } @@ -338,6 +341,10 @@ TOperationId::EKind ParseKind(const std::string_view value) { return TOperationId::FULL_BACKUP; } + if (value.starts_with("analyze")) { + return TOperationId::ANALYZE; + } + return TOperationId::UNUSED; } diff --git a/src/library/operation_id/protos/operation_id.proto b/src/library/operation_id/protos/operation_id.proto index f9c64e04e5b..3791f0d3f08 100644 --- a/src/library/operation_id/protos/operation_id.proto +++ b/src/library/operation_id/protos/operation_id.proto @@ -19,6 +19,7 @@ message TOperationId { RESTORE = 12; COMPACTION = 13; FULL_BACKUP = 14; + ANALYZE = 15; } message TData { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 08c01821a5b..e27fac65f46 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1,3 +1,4 @@ +add_subdirectory(common) add_subdirectory(integration) add_subdirectory(slo_workloads) add_subdirectory(unit) diff --git a/tests/common/CMakeLists.txt b/tests/common/CMakeLists.txt new file mode 100644 index 00000000000..8a198eae24f --- /dev/null +++ b/tests/common/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(iam_mocks) diff --git a/tests/common/iam_mocks/CMakeLists.txt b/tests/common/iam_mocks/CMakeLists.txt new file mode 100644 index 00000000000..ea4490774ac --- /dev/null +++ b/tests/common/iam_mocks/CMakeLists.txt @@ -0,0 +1,15 @@ +_ydb_sdk_add_library(tests-iam-mocks) + +target_link_libraries(tests-iam-mocks PUBLIC + yutil + http-server + http-misc + json + cpp-testing-common + api-client-yc_public +) + +target_sources(tests-iam-mocks PRIVATE + iam_http_mock_server.cpp + iam_grpc_mock_server.cpp +) diff --git a/tests/common/iam_mocks/iam_grpc_mock_server.cpp b/tests/common/iam_mocks/iam_grpc_mock_server.cpp new file mode 100644 index 00000000000..a07a54fef8c --- /dev/null +++ b/tests/common/iam_mocks/iam_grpc_mock_server.cpp @@ -0,0 +1,155 @@ +#include "iam_grpc_mock_server.h" + +#include "iam_test_keys.h" + +namespace NYdb::NTest { + +void TIamTokenServiceStub::SetResponseToken(const std::string& token, int64_t expiresAtSeconds) { + std::lock_guard lock(Lock_); + IamToken_ = token; + ExpiresAtSeconds_ = expiresAtSeconds; +} + +grpc::Status TIamTokenServiceStub::Create( + grpc::ServerContext*, + const yandex::cloud::iam::v1::CreateIamTokenRequest* request, + yandex::cloud::iam::v1::CreateIamTokenResponse* response) +{ + std::lock_guard lock(Lock_); + ++RequestCount_; + LastRequest_ = *request; + HasLastRequest_ = true; + response->set_iam_token(IamToken_); + response->mutable_expires_at()->set_seconds(ExpiresAtSeconds_); + response->mutable_expires_at()->set_nanos(0); + return grpc::Status::OK; +} + +int TIamTokenServiceStub::GetRequestCount() const { + std::lock_guard lock(Lock_); + return RequestCount_; +} + +yandex::cloud::iam::v1::CreateIamTokenRequest TIamTokenServiceStub::GetLastRequest() const { + std::lock_guard lock(Lock_); + return LastRequest_; +} + +bool TIamTokenServiceStub::HasLastRequest() const { + std::lock_guard lock(Lock_); + return HasLastRequest_; +} + +grpc::Status TBlockingIamTokenService::Create( + grpc::ServerContext* context, + const yandex::cloud::iam::v1::CreateIamTokenRequest*, + yandex::cloud::iam::v1::CreateIamTokenResponse* response) +{ + RpcEntered_.store(true); + + std::unique_lock lock(Mutex_); + ReleasedCv_.wait(lock, [this, context] { + return Released_ || ShuttingDown_ || context->IsCancelled(); + }); + + if (context->IsCancelled() || ShuttingDown_) { + return grpc::Status::CANCELLED; + } + response->set_iam_token("released-token"); + response->mutable_expires_at()->set_seconds(4102444800); + response->mutable_expires_at()->set_nanos(0); + return grpc::Status::OK; +} + +void TBlockingIamTokenService::Release() { + { + std::lock_guard lock(Mutex_); + Released_ = true; + } + ReleasedCv_.notify_all(); +} + +void TBlockingIamTokenService::Shutdown() { + { + std::lock_guard lock(Mutex_); + ShuttingDown_ = true; + Released_ = true; + } + ReleasedCv_.notify_all(); +} + +bool TBlockingIamTokenService::WaitUntilRpcEntered(std::chrono::milliseconds timeout) const { + const auto deadline = std::chrono::steady_clock::now() + timeout; + while (std::chrono::steady_clock::now() < deadline) { + if (RpcEntered_.load()) { + return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + return RpcEntered_.load(); +} + +TIamGrpcServer::TIamGrpcServer(grpc::Service* service) + : Service_(service) +{} + +bool TIamGrpcServer::Start() { + grpc::ServerBuilder builder; + int boundPort = 0; + builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &boundPort); + builder.SetMaxSendMessageSize(4 * 1024 * 1024); + builder.RegisterService(Service_); + Server_ = builder.BuildAndStart(); + if (!Server_ || boundPort <= 0) { + return false; + } + Port_ = boundPort; + WaitThread_ = std::thread([this] { Server_->Wait(); }); + return true; +} + +std::string TIamGrpcServer::Endpoint() const { + return "127.0.0.1:" + std::to_string(Port_); +} + +void TIamGrpcServer::Stop() { + if (auto* blocking = dynamic_cast(Service_)) { + blocking->Shutdown(); + } + if (Server_) { + Server_->Shutdown(); + } + if (WaitThread_.joinable()) { + WaitThread_.join(); + } + Server_.reset(); +} + +TIamGrpcServer::~TIamGrpcServer() { + Stop(); +} + +TIamOAuth MakeOAuthParams(const std::string& endpoint, const std::string& oauthToken) { + TIamOAuth params; + params.Endpoint = endpoint; + params.OAuthToken = oauthToken; + params.EnableSsl = false; + params.RefreshPeriod = TDuration::Hours(1); + params.RequestTimeout = TDuration::Seconds(5); + return params; +} + +TIamJwtParams MakeJwtParams(const std::string& endpoint) { + TIamJwtParams params; + params.Endpoint = endpoint; + params.EnableSsl = false; + params.RefreshPeriod = TDuration::Hours(1); + params.RequestTimeout = TDuration::Seconds(5); + params.JwtParams.AccountId = "unit-test-account"; + params.JwtParams.KeyId = "unit-test-key"; + params.JwtParams.PrivKey = TestRSAPrivateKey; + params.JwtParams.PubKey = TestRSAPublicKey; + return params; +} + +} // namespace NYdb::NTest diff --git a/tests/common/iam_mocks/iam_grpc_mock_server.h b/tests/common/iam_mocks/iam_grpc_mock_server.h new file mode 100644 index 00000000000..cb1666636ee --- /dev/null +++ b/tests/common/iam_mocks/iam_grpc_mock_server.h @@ -0,0 +1,100 @@ +#pragma once + +#include + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace NYdb::NTest { + +class TIamTokenServiceStub final : public yandex::cloud::iam::v1::IamTokenService::Service { +public: + void SetResponseToken(const std::string& token, int64_t expiresAtSeconds = 4102444800); + + grpc::Status Create( + grpc::ServerContext*, + const yandex::cloud::iam::v1::CreateIamTokenRequest* request, + yandex::cloud::iam::v1::CreateIamTokenResponse* response) override; + + int GetRequestCount() const; + yandex::cloud::iam::v1::CreateIamTokenRequest GetLastRequest() const; + bool HasLastRequest() const; + +private: + mutable std::mutex Lock_; + std::string IamToken_; + int64_t ExpiresAtSeconds_ = 4102444800; + int RequestCount_ = 0; + yandex::cloud::iam::v1::CreateIamTokenRequest LastRequest_; + bool HasLastRequest_ = false; +}; + +class TBlockingIamTokenService final : public yandex::cloud::iam::v1::IamTokenService::Service { +public: + grpc::Status Create( + grpc::ServerContext* context, + const yandex::cloud::iam::v1::CreateIamTokenRequest*, + yandex::cloud::iam::v1::CreateIamTokenResponse* response) override; + + void Release(); + void Shutdown(); + + bool WaitUntilRpcEntered(std::chrono::milliseconds timeout) const; + +private: + std::atomic RpcEntered_{false}; + bool Released_ = false; + bool ShuttingDown_ = false; + mutable std::mutex Mutex_; + std::condition_variable ReleasedCv_; +}; + +class TBlockingIamReleaseGuard { +public: + explicit TBlockingIamReleaseGuard(TBlockingIamTokenService& service) + : Service_(service) + {} + + ~TBlockingIamReleaseGuard() { + Service_.Release(); + } + +private: + TBlockingIamTokenService& Service_; +}; + +class TIamGrpcServer { +public: + explicit TIamGrpcServer(grpc::Service* service); + TIamGrpcServer(const TIamGrpcServer&) = delete; + TIamGrpcServer& operator=(const TIamGrpcServer&) = delete; + + bool Start(); + std::string Endpoint() const; + void Stop(); + + ~TIamGrpcServer(); + +private: + grpc::Service* Service_ = nullptr; + std::unique_ptr Server_; + int Port_ = 0; + std::thread WaitThread_; +}; + +TIamOAuth MakeOAuthParams(const std::string& endpoint, const std::string& oauthToken = "unit-test-oauth-token"); + +TIamJwtParams MakeJwtParams(const std::string& endpoint); + +} // namespace NYdb::NTest diff --git a/tests/common/iam_mocks/iam_http_assertions.h b/tests/common/iam_mocks/iam_http_assertions.h new file mode 100644 index 00000000000..ddd8220f3f0 --- /dev/null +++ b/tests/common/iam_mocks/iam_http_assertions.h @@ -0,0 +1,18 @@ +#pragma once + +#include "iam_http_mock_server.h" + +#include + +namespace NYdb::NTest { + +inline void AssertMetadataRequestShape(const TMetadataServer& server) { + ASSERT_TRUE(server.HasLastRequest()); + const auto request = server.GetLastRequest(); + EXPECT_TRUE(request.IsValid); + EXPECT_EQ(request.Method, "GET"); + EXPECT_EQ(request.Path, kMetadataTokenPath); + EXPECT_EQ(request.MetadataFlavor, kMetadataFlavorValue); +} + +} // namespace NYdb::NTest diff --git a/tests/common/iam_mocks/iam_http_mock_server.cpp b/tests/common/iam_mocks/iam_http_mock_server.cpp new file mode 100644 index 00000000000..e452c2d493d --- /dev/null +++ b/tests/common/iam_mocks/iam_http_mock_server.cpp @@ -0,0 +1,155 @@ +#include "iam_http_mock_server.h" + +#include +#include + +namespace NYdb::NTest { + +namespace { + +bool IsValidMetadataRequest(const TStringBuf& method, const TStringBuf& path, + const THttpHeaders& headers) +{ + if (method != "GET") { + return false; + } + if (path != kMetadataTokenPath) { + return false; + } + const THttpInputHeader* flavorHeader = headers.FindHeader(kMetadataFlavorHeader); + return flavorHeader && flavorHeader->Value() == kMetadataFlavorValue; +} + +} // namespace + +bool TMetadataServer::TRequest::DoReply(const TReplyParams& params) { + const TParsedHttpFull parsed(params.Input.FirstLine()); + const auto& headers = params.Input.Headers(); + + TMetadataRequestInfo requestInfo; + requestInfo.Method = TString{parsed.Method}; + requestInfo.Path = TString{parsed.Path}; + if (const THttpInputHeader* flavorHeader = headers.FindHeader(kMetadataFlavorHeader)) { + requestInfo.MetadataFlavor = flavorHeader->Value(); + } + requestInfo.IsValid = IsValidMetadataRequest(parsed.Method, parsed.Path, headers); + + HttpCodes statusCode; + std::string responseBody; + bool strictMode; + { + std::lock_guard lock(Server->Lock_); + ++Server->RequestCount_; + Server->LastRequest_ = std::move(requestInfo); + Server->HasLastRequest_ = true; + strictMode = Server->StrictMode_; + if (!Server->LastRequest_.IsValid && strictMode) { + statusCode = HTTP_NOT_FOUND; + responseBody = ""; + } else { + statusCode = Server->StatusCode_; + responseBody = Server->Response_; + } + } + + THttpResponse resp(statusCode); + resp.SetContent(TString{responseBody}); + resp.OutTo(params.Output); + return true; +} + +TMetadataServer::TMetadataServer() + : PortHolder(NTesting::GetFreePort()) + , Port(static_cast(PortHolder)) + , HttpOptions(Port) + , HttpServer(this, HttpOptions) +{ + HttpServer.Start(); +} + +TMetadataServer::~TMetadataServer() { + HttpServer.Stop(); +} + +TClientRequest* TMetadataServer::CreateClient() { + return new TRequest(this); +} + +void TMetadataServer::SetResponse(HttpCodes code, const std::string& response) { + std::lock_guard lock(Lock_); + StatusCode_ = code; + Response_ = response; +} + +void TMetadataServer::SetStrictMode(bool strict) { + std::lock_guard lock(Lock_); + StrictMode_ = strict; +} + +int TMetadataServer::GetRequestCount() const { + std::lock_guard lock(Lock_); + return RequestCount_; +} + +TMetadataRequestInfo TMetadataServer::GetLastRequest() const { + std::lock_guard lock(Lock_); + return LastRequest_; +} + +bool TMetadataServer::HasLastRequest() const { + std::lock_guard lock(Lock_); + return HasLastRequest_; +} + +std::string MakeTokenResponse(const std::string& token, int expiresIn) { + TStringStream ss; + NJson::TJsonWriter w(&ss, false); + w.OpenMap(); + w.WriteKey("access_token"); + w.Write(token); + w.WriteKey("token_type"); + w.Write("Bearer"); + w.WriteKey("expires_in"); + w.Write(expiresIn); + w.CloseMap(); + w.Flush(); + return ss.Str(); +} + +std::string MakeTokenResponseWithExpiry(const std::string& token, const TInstant& expiry) { + TStringStream ss; + NJson::TJsonWriter w(&ss, false); + w.OpenMap(); + w.WriteKey("access_token"); + w.Write(token); + w.WriteKey("token_type"); + w.Write("Bearer"); + w.WriteKey("expiry"); + w.Write(expiry.FormatGmTime("%Y-%m-%dT%H:%M:%S.000000000Z")); + w.CloseMap(); + w.Flush(); + return ss.Str(); +} + +std::string MakeTokenResponseNoExpiry(const std::string& token) { + TStringStream ss; + NJson::TJsonWriter w(&ss, false); + w.OpenMap(); + w.WriteKey("access_token"); + w.Write(token); + w.WriteKey("token_type"); + w.Write("Bearer"); + w.CloseMap(); + w.Flush(); + return ss.Str(); +} + +TIamHost MakeMetadataParams(uint16_t port) { + TIamHost params; + params.Host = "localhost"; + params.Port = port; + params.RefreshPeriod = TDuration::Hours(1); + return params; +} + +} // namespace NYdb::NTest diff --git a/tests/common/iam_mocks/iam_http_mock_server.h b/tests/common/iam_mocks/iam_http_mock_server.h new file mode 100644 index 00000000000..1d0aee5d5b8 --- /dev/null +++ b/tests/common/iam_mocks/iam_http_mock_server.h @@ -0,0 +1,74 @@ +#pragma once + +#include + +#include +#include +#include + +#include + +#include +#include + +namespace NYdb::NTest { + +inline constexpr const char* kMetadataTokenPath = + "/computeMetadata/v1/instance/service-accounts/default/token"; +inline constexpr const char* kMetadataFlavorHeader = "Metadata-Flavor"; +inline constexpr const char* kMetadataFlavorValue = "Google"; + +struct TMetadataRequestInfo { + std::string Method; + std::string Path; + std::string MetadataFlavor; + bool IsValid = false; +}; + +class TMetadataServer : public THttpServer::ICallBack { +public: + class TRequest : public TRequestReplier { + public: + explicit TRequest(TMetadataServer* server) + : Server(server) + {} + + bool DoReply(const TReplyParams& params) override; + + TMetadataServer* Server = nullptr; + }; + + TMetadataServer(); + ~TMetadataServer(); + + TClientRequest* CreateClient() override; + + void SetResponse(HttpCodes code, const std::string& response); + void SetStrictMode(bool strict); + + int GetRequestCount() const; + TMetadataRequestInfo GetLastRequest() const; + bool HasLastRequest() const; + + NTesting::TPortHolder PortHolder; + uint16_t Port; + THttpServer::TOptions HttpOptions; + THttpServer HttpServer; + +private: + mutable std::mutex Lock_; + HttpCodes StatusCode_ = HTTP_OK; + std::string Response_; + bool StrictMode_ = true; + int RequestCount_ = 0; + TMetadataRequestInfo LastRequest_; + bool HasLastRequest_ = false; +}; + +std::string MakeTokenResponse(const std::string& token, int expiresIn); +std::string MakeTokenResponseWithExpiry(const std::string& token, const TInstant& expiry); +std::string MakeTokenResponseNoExpiry(const std::string& token); + +TIamHost MakeMetadataParams(uint16_t port); + +} // namespace NYdb::NTest diff --git a/tests/common/iam_mocks/iam_jwt_assertions.h b/tests/common/iam_mocks/iam_jwt_assertions.h new file mode 100644 index 00000000000..5a4c272b950 --- /dev/null +++ b/tests/common/iam_mocks/iam_jwt_assertions.h @@ -0,0 +1,56 @@ +#pragma once + +#include "iam_test_keys.h" + +#include + +#include + +#include +#include +#include + +namespace NYdb::NTest { + +inline void AssertIamJwt(const std::string& jwt) { + ASSERT_FALSE(jwt.empty()); + +#ifdef YDB_SDK_OSS + using TDecodedJwt = jwt::decoded_jwt; +#else + using TDecodedJwt = jwt::decoded_jwt; +#endif + + const TDecodedJwt decoded(jwt); + EXPECT_EQ(decoded.get_algorithm(), "PS256"); + + const std::string data = decoded.get_header_base64() + "." + decoded.get_payload_base64(); + const std::string signature = decoded.get_signature(); + jwt::algorithm::ps256 verifier(TestRSAPublicKey, ""); +#ifdef YDB_SDK_OSS + std::error_code ec; + verifier.verify(data, signature, ec); + ASSERT_FALSE(ec) << ec.message(); +#else + verifier.verify(data, signature); +#endif + + ASSERT_TRUE(decoded.has_key_id()); + EXPECT_EQ(decoded.get_key_id(), kIamJwtKeyId); + + ASSERT_TRUE(decoded.has_issuer()); + EXPECT_EQ(decoded.get_issuer(), kIamJwtIssuer); + + ASSERT_TRUE(decoded.has_audience()); + EXPECT_EQ(decoded.get_audience(), std::set{kIamJwtAudience}); + + const auto now = std::chrono::system_clock::now(); + const auto iat = decoded.get_issued_at(); + const auto exp = decoded.get_expires_at(); + EXPECT_GE(iat, now - std::chrono::minutes(10)); + EXPECT_LE(iat, now + std::chrono::minutes(1)); + EXPECT_GT(exp, iat); + EXPECT_LE(exp - iat, std::chrono::hours(2)); +} + +} // namespace NYdb::NTest diff --git a/tests/common/iam_mocks/iam_test_keys.h b/tests/common/iam_mocks/iam_test_keys.h new file mode 100644 index 00000000000..191c65ce38e --- /dev/null +++ b/tests/common/iam_mocks/iam_test_keys.h @@ -0,0 +1,74 @@ +#pragma once + +#include + +#include + +#include + +namespace NYdb::NTest { + +// PS256-compatible RSA test keypair (also used in jwt_token_source_ut.cpp). +inline constexpr const char* TestRSAPrivateKey = + "-----BEGIN PRIVATE KEY-----\n" + "MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQC75/JS3rMcLJxv\n" + "FgpOzF5+2gH+Yig3RE2MTl9uwC0BZKAv6foYr7xywQyWIK+W1cBhz8R4LfFmZo2j\n" + "M0aCvdRmNBdW0EDSTnHLxCsFhoQWLVq+bI5f5jzkcoiioUtaEpADPqwgVULVtN/n\n" + "nPJiZ6/dU30C3jmR6+LUgEntUtWt3eq3xQIn5lG3zC1klBY/HxtfH5Hu8xBvwRQT\n" + "Jnh3UpPLj8XwSmriDgdrhR7o6umWyVuGrMKlLHmeivlfzjYtfzO1MOIMG8t2/zxG\n" + "R+xb4Vwks73sH1KruH/0/JMXU97npwpe+Um+uXhpldPygGErEia7abyZB2gMpXqr\n" + "WYKMo02NAgMBAAECggEAO0BpC5OYw/4XN/optu4/r91bupTGHKNHlsIR2rDzoBhU\n" + "YLd1evpTQJY6O07EP5pYZx9mUwUdtU4KRJeDGO/1/WJYp7HUdtxwirHpZP0lQn77\n" + "uccuX/QQaHLrPekBgz4ONk+5ZBqukAfQgM7fKYOLk41jgpeDbM2Ggb6QUSsJISEp\n" + "zrwpI/nNT/wn+Hvx4DxrzWU6wF+P8kl77UwPYlTA7GsT+T7eKGVH8xsxmK8pt6lg\n" + "svlBA5XosWBWUCGLgcBkAY5e4ZWbkdd183o+oMo78id6C+PQPE66PLDtHWfpRRmN\n" + "m6XC03x6NVhnfvfozoWnmS4+e4qj4F/emCHvn0GMywKBgQDLXlj7YPFVXxZpUvg/\n" + "rheVcCTGbNmQJ+4cZXx87huqwqKgkmtOyeWsRc7zYInYgraDrtCuDBCfP//ZzOh0\n" + "LxepYLTPk5eNn/GT+VVrqsy35Ccr60g7Lp/bzb1WxyhcLbo0KX7/6jl0lP+VKtdv\n" + "mto+4mbSBXSM1Y5BVVoVgJ3T/wKBgQDsiSvPRzVi5TTj13x67PFymTMx3HCe2WzH\n" + "JUyepCmVhTm482zW95pv6raDr5CTO6OYpHtc5sTTRhVYEZoEYFTM9Vw8faBtluWG\n" + "BjkRh4cIpoIARMn74YZKj0C/0vdX7SHdyBOU3bgRPHg08Hwu3xReqT1kEPSI/B2V\n" + "4pe5fVrucwKBgQCNFgUxUA3dJjyMES18MDDYUZaRug4tfiYouRdmLGIxUxozv6CG\n" + "ZnbZzwxFt+GpvPUV4f+P33rgoCvFU+yoPctyjE6j+0aW0DFucPmb2kBwCu5J/856\n" + "kFwCx3blbwFHAco+SdN7g2kcwgmV2MTg/lMOcU7XwUUcN0Obe7UlWbckzQKBgQDQ\n" + "nXaXHL24GGFaZe4y2JFmujmNy1dEsoye44W9ERpf9h1fwsoGmmCKPp90az5+rIXw\n" + "FXl8CUgk8lXW08db/r4r+ma8Lyx0GzcZyplAnaB5/6j+pazjSxfO4KOBy4Y89Tb+\n" + "TP0AOcCi6ws13bgY+sUTa/5qKA4UVw+c5zlb7nRpgwKBgGXAXhenFw1666482iiN\n" + "cHSgwc4ZHa1oL6aNJR1XWH+aboBSwR+feKHUPeT4jHgzRGo/aCNHD2FE5I8eBv33\n" + "of1kWYjAO0YdzeKrW0rTwfvt9gGg+CS397aWu4cy+mTI+MNfBgeDAIVBeJOJXLlX\n" + "hL8bFAuNNVrCOp79TNnNIsh7\n" + "-----END PRIVATE KEY-----\n"; + +inline constexpr const char* TestRSAPublicKey = + "-----BEGIN PUBLIC KEY-----\n" + "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAu+fyUt6zHCycbxYKTsxe\n" + "ftoB/mIoN0RNjE5fbsAtAWSgL+n6GK+8csEMliCvltXAYc/EeC3xZmaNozNGgr3U\n" + "ZjQXVtBA0k5xy8QrBYaEFi1avmyOX+Y85HKIoqFLWhKQAz6sIFVC1bTf55zyYmev\n" + "3VN9At45kevi1IBJ7VLVrd3qt8UCJ+ZRt8wtZJQWPx8bXx+R7vMQb8EUEyZ4d1KT\n" + "y4/F8Epq4g4Ha4Ue6OrplslbhqzCpSx5nor5X842LX8ztTDiDBvLdv88RkfsW+Fc\n" + "JLO97B9Sq7h/9PyTF1Pe56cKXvlJvrl4aZXT8oBhKxImu2m8mQdoDKV6q1mCjKNN\n" + "jQIDAQAB\n" + "-----END PUBLIC KEY-----\n"; + +inline constexpr const char* kIamJwtAudience = "https://iam.api.cloud.yandex.net/iam/v1/tokens"; +inline constexpr const char* kIamJwtKeyId = "unit-test-key"; +inline constexpr const char* kIamJwtIssuer = "unit-test-account"; + +inline std::string MakeJwtKeyFileContent() { + TStringStream ss; + NJson::TJsonWriter w(&ss, false); + w.OpenMap(); + w.WriteKey("id"); + w.Write(kIamJwtKeyId); + w.WriteKey("service_account_id"); + w.Write(kIamJwtIssuer); + w.WriteKey("private_key"); + w.Write(TestRSAPrivateKey); + w.WriteKey("public_key"); + w.Write(TestRSAPublicKey); + w.CloseMap(); + w.Flush(); + return ss.Str(); +} + +} // namespace NYdb::NTest diff --git a/tests/integration/CMakeLists.txt b/tests/integration/CMakeLists.txt index 8a4d0d58d61..3bcef103686 100644 --- a/tests/integration/CMakeLists.txt +++ b/tests/integration/CMakeLists.txt @@ -1,4 +1,5 @@ add_subdirectory(auth) +add_subdirectory(iam) add_subdirectory(basic_example) add_subdirectory(bulk_upsert) add_subdirectory(key_conflict) diff --git a/tests/integration/bulk_upsert/bulk_upsert.cpp b/tests/integration/bulk_upsert/bulk_upsert.cpp index de419179fbe..2a3e1083aa6 100644 --- a/tests/integration/bulk_upsert/bulk_upsert.cpp +++ b/tests/integration/bulk_upsert/bulk_upsert.cpp @@ -3,8 +3,6 @@ #include #include -static constexpr size_t BATCH_SIZE = 1000; - TRunArgs GetRunArgs() { std::string endpoint = std::getenv("YDB_ENDPOINT"); std::string database = std::getenv("YDB_DATABASE"); @@ -37,13 +35,15 @@ TStatus CreateTable(TTableClient& client, const std::string& table) { return status; } -TStatistic GetLogBatch(uint64_t logOffset, std::vector& logBatch, uint32_t lastNumber) { +TStatistic GetLogBatch(uint64_t logOffset, std::vector& logBatch, uint32_t lastNumber, + size_t batchSize) +{ logBatch.clear(); uint32_t correctSumApp = 0; uint32_t correctSumHost = 0; uint32_t correctRowCount = 0; - for (size_t i = 0; i < BATCH_SIZE; ++i) { + for (size_t i = 0; i < batchSize; ++i) { TLogMessage message; message.Pk.Id = correctRowCount + lastNumber; message.Pk.App = "App_" + std::to_string(logOffset % 10); @@ -61,8 +61,7 @@ TStatistic GetLogBatch(uint64_t logOffset, std::vector& logBatch, u return {correctSumApp, correctSumHost, correctRowCount}; } -TStatus WriteLogBatch(TTableClient& tableClient, const std::string& table, const std::vector& logBatch, - const TRetryOperationSettings& retrySettings) { +TValue BuildLogBatchValue(const std::vector& logBatch) { TValueBuilder rows; rows.BeginList(); for (const auto& message : logBatch) { @@ -77,7 +76,36 @@ TStatus WriteLogBatch(TTableClient& tableClient, const std::string& table, const .EndStruct(); } rows.EndList(); - auto bulkUpsertOperation = [table, rowsValue = rows.Build()](TTableClient& tableClient) { + return rows.Build(); +} + +TDuration MeasureBulkUpsertWallTime( + TTableClient& tableClient, + const std::string& table, + uint64_t logOffset, + uint32_t idOffset, + const TBulkUpsertSettings& settings, + size_t batchSize) +{ + std::vector logBatch; + GetLogBatch(logOffset, logBatch, idOffset, batchSize); + auto rows = BuildLogBatchValue(logBatch); + + const auto start = TInstant::Now(); + const auto status = tableClient.BulkUpsert(table, std::move(rows), settings).GetValueSync(); + const auto elapsed = TInstant::Now() - start; + + if (!status.IsSuccess()) { + throw NStatusHelpers::TYdbErrorException(status); + } + + return elapsed; +} + +TStatus WriteLogBatch(TTableClient& tableClient, const std::string& table, const std::vector& logBatch, + const TRetryOperationSettings& retrySettings) { + auto rowsValue = BuildLogBatchValue(logBatch); + auto bulkUpsertOperation = [table, rowsValue = std::move(rowsValue)](TTableClient& tableClient) { TValue r = rowsValue; auto status = tableClient.BulkUpsert(table, std::move(r)); return status.GetValueSync(); diff --git a/tests/integration/bulk_upsert/bulk_upsert.h b/tests/integration/bulk_upsert/bulk_upsert.h index e6d13a5d7f7..eb35ed2a909 100644 --- a/tests/integration/bulk_upsert/bulk_upsert.h +++ b/tests/integration/bulk_upsert/bulk_upsert.h @@ -6,6 +6,8 @@ using namespace NYdb; using namespace NYdb::NTable; +constexpr size_t BULK_UPSERT_BATCH_SIZE = 1000; + struct TRunArgs { TDriver Driver; std::string Path; @@ -32,8 +34,17 @@ struct TStatistic { TRunArgs GetRunArgs(); TStatus CreateTable(TTableClient& client, const std::string& table); -TStatistic GetLogBatch(uint64_t logOffset, std::vector& logBatch, uint32_t lastNumber); +TStatistic GetLogBatch(uint64_t logOffset, std::vector& logBatch, uint32_t lastNumber, + size_t batchSize = BULK_UPSERT_BATCH_SIZE); +TValue BuildLogBatchValue(const std::vector& logBatch); TStatus WriteLogBatch(TTableClient& tableClient, const std::string& table, const std::vector& logBatch, const TRetryOperationSettings& retrySettings); +TDuration MeasureBulkUpsertWallTime( + TTableClient& tableClient, + const std::string& table, + uint64_t logOffset, + uint32_t idOffset, + const TBulkUpsertSettings& settings, + size_t batchSize = BULK_UPSERT_BATCH_SIZE); TStatistic Select(TTableClient& client, const std::string& path); void DropTable(TTableClient& client, const std::string& path); diff --git a/tests/integration/bulk_upsert/main.cpp b/tests/integration/bulk_upsert/main.cpp index 5dcfd0f2a23..871e4631839 100644 --- a/tests/integration/bulk_upsert/main.cpp +++ b/tests/integration/bulk_upsert/main.cpp @@ -51,3 +51,68 @@ TEST(BulkUpsert, BulkUpsert) { DropTable(client, path); driver.Stop(true); } + +TEST(BulkUpsert, RetryOverheadOnHappyPath) { + constexpr size_t kPerfBatchSize = 10000; + constexpr size_t kWarmupIterations = 3; + constexpr size_t kMeasuredIterations = 30; + // Retry wrapper adds client-side bookkeeping; server RPC dominates, so allow modest variance. + constexpr double kMaxSlowdownRatio = 1.1; + + auto [driver, basePath] = GetRunArgs(); + const std::string tableNoRetry = basePath + "_perf_noretry"; + const std::string tableWithRetry = basePath + "_perf_retry"; + + TTableClient client(driver); + + const auto statusCreateNoRetry = CreateTable(client, tableNoRetry); + ASSERT_TRUE(statusCreateNoRetry.IsSuccess()) << ToString(statusCreateNoRetry); + const auto statusCreateWithRetry = CreateTable(client, tableWithRetry); + ASSERT_TRUE(statusCreateWithRetry.IsSuccess()) << ToString(statusCreateWithRetry); + + const TBulkUpsertSettings noRetrySettings = [] { + TBulkUpsertSettings settings; + settings.RetrySettings(TRetryOperationSettings().MaxRetries(0)); + return settings; + }(); + + const TBulkUpsertSettings withRetrySettings = [] { + TBulkUpsertSettings settings; + settings.RetrySettings(TRetryOperationSettings().MaxRetries(10).Idempotent(true)); + return settings; + }(); + + auto runInterleavedIterations = [&](size_t iterations) { + TDuration noRetryTotal; + TDuration withRetryTotal; + for (size_t i = 0; i < iterations; ++i) { + const uint32_t offset = static_cast(i * kPerfBatchSize); + noRetryTotal += MeasureBulkUpsertWallTime( + client, tableNoRetry, i, offset, noRetrySettings, kPerfBatchSize); + withRetryTotal += MeasureBulkUpsertWallTime( + client, tableWithRetry, i, offset, withRetrySettings, kPerfBatchSize); + } + return std::make_pair(noRetryTotal, withRetryTotal); + }; + + try { + runInterleavedIterations(kWarmupIterations); + + const auto [noRetryTime, withRetryTime] = runInterleavedIterations(kMeasuredIterations); + + const double ratio = static_cast(withRetryTime.MicroSeconds()) + / static_cast(noRetryTime.MicroSeconds()); + + EXPECT_LE(ratio, kMaxSlowdownRatio) + << "BulkUpsert with retries should not be significantly slower on the happy path." + << " noRetryTotalUs=" << noRetryTime.MicroSeconds() + << " withRetryTotalUs=" << withRetryTime.MicroSeconds() + << " ratio=" << ratio; + } catch (const NYdb::NStatusHelpers::TYdbErrorException& e) { + FAIL() << "BulkUpsert benchmark failed:\n" << e.what(); + } + + DropTable(client, tableNoRetry); + DropTable(client, tableWithRetry); + driver.Stop(true); +} diff --git a/tests/integration/iam/CMakeLists.txt b/tests/integration/iam/CMakeLists.txt new file mode 100644 index 00000000000..6ecaf73d5ab --- /dev/null +++ b/tests/integration/iam/CMakeLists.txt @@ -0,0 +1,47 @@ +add_ydb_test(NAME iam_metadata_it GTEST + SOURCES + iam_test_fixture.cpp + metadata_it.cpp + LINK_LIBRARIES + tests-iam-mocks + YDB-CPP-SDK::Iam + YDB-CPP-SDK::Query + LABELS + integration +) + +add_ydb_test(NAME iam_driver_auth_it GTEST + SOURCES + iam_test_fixture.cpp + driver_auth_it.cpp + LINK_LIBRARIES + tests-iam-mocks + YDB-CPP-SDK::Iam + YDB-CPP-SDK::Query + LABELS + integration +) + +add_ydb_test(NAME iam_jwt_it GTEST + SOURCES + iam_test_fixture.cpp + jwt_it.cpp + LINK_LIBRARIES + tests-iam-mocks + YDB-CPP-SDK::Iam + YDB-CPP-SDK::Query + LABELS + integration +) + +add_ydb_test(NAME iam_oauth_it GTEST + SOURCES + iam_test_fixture.cpp + oauth_it.cpp + LINK_LIBRARIES + tests-iam-mocks + YDB-CPP-SDK::Iam + YDB-CPP-SDK::Query + LABELS + integration +) diff --git a/tests/integration/iam/driver_auth_it.cpp b/tests/integration/iam/driver_auth_it.cpp new file mode 100644 index 00000000000..1cfa4f5ad18 --- /dev/null +++ b/tests/integration/iam/driver_auth_it.cpp @@ -0,0 +1,35 @@ +// Auth-enforcement negative control only; positive driver wiring lives in oauth/jwt/metadata *_it.cpp. +#include +#include +#include +#include +#include + +#include + +using namespace NYdb; +using namespace NYdb::NTest; + +namespace { + +constexpr const char* kInvalidToken = "invalid-e2e-token"; + +} // namespace + +TEST(DriverAuth, InvalidStaticTokenRejected) { + AssertAuthFailure( + CreateOAuthCredentialsProviderFactory(kInvalidToken), + "invalid static token"); +} + +TEST(DriverAuth, MetadataInvalidTokenRejected) { + TMetadataServer server; + server.SetResponse(HTTP_OK, MakeTokenResponse(kInvalidToken, 3600)); + + AssertAuthFailure( + CreateIamCredentialsProviderFactory(MakeMetadataParams(server.Port)), + "invalid metadata token"); + + EXPECT_GE(server.GetRequestCount(), 1); + AssertMetadataRequestShape(server); +} diff --git a/tests/integration/iam/iam_test_fixture.cpp b/tests/integration/iam/iam_test_fixture.cpp new file mode 100644 index 00000000000..e311d1128f1 --- /dev/null +++ b/tests/integration/iam/iam_test_fixture.cpp @@ -0,0 +1,68 @@ +#include "iam_test_fixture.h" + +#include + +#include + +#include +#include +#include + +namespace NYdb::NTest { + +namespace { + +struct TYdbEnv { + std::string Endpoint; + std::string Database; +}; + +TYdbEnv RequireYdbEnvImpl() { + const char* endpoint = std::getenv("YDB_ENDPOINT"); + if (!endpoint || !*endpoint) { + throw std::runtime_error( + "YDB_ENDPOINT is not set; recipe-backed integration target must provide it"); + } + const char* database = std::getenv("YDB_DATABASE"); + if (!database || !*database) { + throw std::runtime_error( + "YDB_DATABASE is not set; recipe-backed integration target must provide it"); + } + return TYdbEnv{endpoint, database}; +} + +} // namespace + +TDriverConfig MakeDriverConfig(TCredentialsProviderFactoryPtr factory) { + const TYdbEnv env = RequireYdbEnvImpl(); + return TDriverConfig() + .SetEndpoint(env.Endpoint) + .SetDatabase(env.Database) + .SetCredentialsProviderFactory(std::move(factory)); +} + +TStatus RunSelect1Status(TDriver& driver) { + NQuery::TQueryClient client(driver); + return client.ExecuteQuery("SELECT 1", NQuery::TTxControl::NoTx()).GetValueSync(); +} + +void RunSelect1ExpectSuccess(TDriver& driver) { + const auto status = RunSelect1Status(driver); + ASSERT_TRUE(status.IsSuccess()) << status.GetIssues().ToString(); +} + +bool IsAuthError(const TStatus& status) { + const EStatus code = status.GetStatus(); + return code == EStatus::CLIENT_UNAUTHENTICATED || code == EStatus::UNAUTHORIZED; +} + +void AssertAuthFailure(TCredentialsProviderFactoryPtr factory, const char* context) { + TDriver driver(MakeDriverConfig(std::move(factory))); + const auto status = RunSelect1Status(driver); + ASSERT_FALSE(status.IsSuccess()) + << "Expected auth failure with " << context << ", but query succeeded. " + << "Recipe must set YDB_ENFORCE_USER_TOKEN_REQUIREMENT=true."; + EXPECT_TRUE(IsAuthError(status)) << status.GetIssues().ToString(); +} + +} // namespace NYdb::NTest diff --git a/tests/integration/iam/iam_test_fixture.h b/tests/integration/iam/iam_test_fixture.h new file mode 100644 index 00000000000..48c9e786adc --- /dev/null +++ b/tests/integration/iam/iam_test_fixture.h @@ -0,0 +1,19 @@ +#pragma once + +#include +#include +#include + +namespace NYdb::NTest { + +inline constexpr const char* kMockRootBuiltinToken = "root@builtin"; + +TDriverConfig MakeDriverConfig(TCredentialsProviderFactoryPtr factory); + +TStatus RunSelect1Status(TDriver& driver); +void RunSelect1ExpectSuccess(TDriver& driver); + +bool IsAuthError(const TStatus& status); +void AssertAuthFailure(TCredentialsProviderFactoryPtr factory, const char* context); + +} // namespace NYdb::NTest diff --git a/tests/integration/iam/jwt_it.cpp b/tests/integration/iam/jwt_it.cpp new file mode 100644 index 00000000000..f276e3134b6 --- /dev/null +++ b/tests/integration/iam/jwt_it.cpp @@ -0,0 +1,99 @@ +#include +#include +#include +#include +#include + +#include + +#include +#include + +using namespace NYdb; +using namespace NYdb::NTest; + +namespace { + +class TJwtIamFixture : public ::testing::Test { +protected: + void SetUp() override { + Stub_.SetResponseToken(kMockRootBuiltinToken); + ASSERT_TRUE(Server_.Start()); + } + + TIamJwtContent MakeJwtContentParams() const { + TIamJwtContent params; + params.Endpoint = Server_.Endpoint(); + params.EnableSsl = false; + params.RefreshPeriod = TDuration::Hours(1); + params.RequestTimeout = TDuration::Seconds(5); + params.JwtContent = MakeJwtKeyFileContent(); + return params; + } + + TIamJwtFilename MakeJwtFileParams(const TString& keyPath) const { + TIamJwtFilename params; + params.Endpoint = Server_.Endpoint(); + params.EnableSsl = false; + params.RefreshPeriod = TDuration::Hours(1); + params.RequestTimeout = TDuration::Seconds(5); + params.JwtFilename = keyPath; + return params; + } + + TIamTokenServiceStub Stub_; + TIamGrpcServer Server_{&Stub_}; +}; + +TString WriteJwtKeyToTempFile(TTempDir& tempDir) { + const TString keyPath = tempDir.Path() / "sa-key.json"; + TFileOutput out(keyPath); + out.Write(MakeJwtKeyFileContent()); + return keyPath; +} + +} // namespace + +TEST_F(TJwtIamFixture, JwtParams_NoArgCreateProvider) { + auto factory = CreateIamJwtParamsCredentialsProviderFactory(MakeJwtContentParams()); + auto provider = factory->CreateProvider(); + + EXPECT_EQ(provider->GetAuthInfo(), kMockRootBuiltinToken); + ASSERT_TRUE(Stub_.HasLastRequest()); + AssertIamJwt(Stub_.GetLastRequest().jwt()); +} + +TEST_F(TJwtIamFixture, JwtParams_DriverUsesMockIamToken) { + auto factory = CreateIamJwtParamsCredentialsProviderFactory(MakeJwtContentParams()); + TDriver driver(MakeDriverConfig(factory)); + RunSelect1ExpectSuccess(driver); + + EXPECT_GE(Stub_.GetRequestCount(), 1); + ASSERT_TRUE(Stub_.HasLastRequest()); + AssertIamJwt(Stub_.GetLastRequest().jwt()); +} + +TEST_F(TJwtIamFixture, JwtFile_NoArgCreateProvider) { + TTempDir tempDir; + const TString keyPath = WriteJwtKeyToTempFile(tempDir); + + auto factory = CreateIamJwtFileCredentialsProviderFactory(MakeJwtFileParams(keyPath)); + auto provider = factory->CreateProvider(); + + EXPECT_EQ(provider->GetAuthInfo(), kMockRootBuiltinToken); + ASSERT_TRUE(Stub_.HasLastRequest()); + AssertIamJwt(Stub_.GetLastRequest().jwt()); +} + +TEST_F(TJwtIamFixture, JwtFile_DriverUsesMockIamToken) { + TTempDir tempDir; + const TString keyPath = WriteJwtKeyToTempFile(tempDir); + + auto factory = CreateIamJwtFileCredentialsProviderFactory(MakeJwtFileParams(keyPath)); + TDriver driver(MakeDriverConfig(factory)); + RunSelect1ExpectSuccess(driver); + + EXPECT_GE(Stub_.GetRequestCount(), 1); + ASSERT_TRUE(Stub_.HasLastRequest()); + AssertIamJwt(Stub_.GetLastRequest().jwt()); +} diff --git a/tests/integration/iam/metadata_it.cpp b/tests/integration/iam/metadata_it.cpp new file mode 100644 index 00000000000..8c0bdb45777 --- /dev/null +++ b/tests/integration/iam/metadata_it.cpp @@ -0,0 +1,69 @@ +#include +#include +#include +#include + +#include + +#include + +using namespace NYdb; +using namespace NYdb::NTest; + +namespace { + +class TMetadataFixture : public ::testing::Test { +protected: + void SetUp() override { + Server_.SetResponse(HTTP_OK, MakeTokenResponse(kMockRootBuiltinToken, 3600)); + } + + TMetadataServer Server_; +}; + +} // namespace + +TEST_F(TMetadataFixture, Metadata_NoArgCreateProvider) { + auto factory = CreateIamCredentialsProviderFactory(MakeMetadataParams(Server_.Port)); + auto provider = factory->CreateProvider(); + + EXPECT_EQ(provider->GetAuthInfo(), kMockRootBuiltinToken); + EXPECT_GE(Server_.GetRequestCount(), 1); + AssertMetadataRequestShape(Server_); +} + +TEST_F(TMetadataFixture, Metadata_ViaFacilityDefault) { + auto factory = CreateIamCredentialsProviderFactory(MakeMetadataParams(Server_.Port)); + auto facility = std::make_shared(); + auto provider = factory->CreateProvider(facility); + + EXPECT_EQ(provider->GetAuthInfo(), kMockRootBuiltinToken); + EXPECT_GE(Server_.GetRequestCount(), 1); + AssertMetadataRequestShape(Server_); +} + +TEST_F(TMetadataFixture, Metadata_DriverUsesMockIamToken) { + auto factory = CreateIamCredentialsProviderFactory(MakeMetadataParams(Server_.Port)); + TDriver driver(MakeDriverConfig(factory)); + RunSelect1ExpectSuccess(driver); + + EXPECT_GE(Server_.GetRequestCount(), 1); + AssertMetadataRequestShape(Server_); +} + +TEST_F(TMetadataFixture, Metadata_PreflightThenDriver) { + auto factory = CreateIamCredentialsProviderFactory(MakeMetadataParams(Server_.Port)); + + // Mirrors CreateFromEnvironment metadata preflight in helpers.cpp. + auto preflightProvider = factory->CreateProvider(); + EXPECT_EQ(preflightProvider->GetAuthInfo(), kMockRootBuiltinToken); + + const int countAfterPreflight = Server_.GetRequestCount(); + AssertMetadataRequestShape(Server_); + + TDriver driver(MakeDriverConfig(factory)); + RunSelect1ExpectSuccess(driver); + + EXPECT_GT(Server_.GetRequestCount(), countAfterPreflight); + AssertMetadataRequestShape(Server_); +} diff --git a/tests/integration/iam/oauth_it.cpp b/tests/integration/iam/oauth_it.cpp new file mode 100644 index 00000000000..0a9b3afc98b --- /dev/null +++ b/tests/integration/iam/oauth_it.cpp @@ -0,0 +1,77 @@ +#include +#include +#include + +#include + +#include + +using namespace NYdb; +using namespace NYdb::NTest; + +namespace { + +constexpr const char* kOAuthToken = "unit-test-oauth-token"; + +class TOAuthIamFixture : public ::testing::Test { +protected: + void SetUp() override { + Stub_.SetResponseToken(kMockRootBuiltinToken); + ASSERT_TRUE(Server_.Start()); + } + + TIamTokenServiceStub Stub_; + TIamGrpcServer Server_{&Stub_}; +}; + +} // namespace + +TEST_F(TOAuthIamFixture, OAuth_NoArgCreateProvider) { + auto factory = CreateIamOAuthCredentialsProviderFactory( + MakeOAuthParams(Server_.Endpoint(), kOAuthToken)); + + auto provider = factory->CreateProvider(); + EXPECT_EQ(provider->GetAuthInfo(), kMockRootBuiltinToken); + ASSERT_TRUE(Stub_.HasLastRequest()); + EXPECT_EQ(Stub_.GetLastRequest().yandex_passport_oauth_token(), kOAuthToken); +} + +TEST_F(TOAuthIamFixture, OAuth_DriverUsesMockIamToken) { + auto factory = CreateIamOAuthCredentialsProviderFactory( + MakeOAuthParams(Server_.Endpoint(), kOAuthToken)); + + TDriver driver(MakeDriverConfig(factory)); + RunSelect1ExpectSuccess(driver); + + EXPECT_GE(Stub_.GetRequestCount(), 1); + ASSERT_TRUE(Stub_.HasLastRequest()); + EXPECT_EQ(Stub_.GetLastRequest().yandex_passport_oauth_token(), kOAuthToken); +} + +TEST_F(TOAuthIamFixture, OAuth_RefreshUsesCachedToken) { + auto factory = CreateIamOAuthCredentialsProviderFactory( + MakeOAuthParams(Server_.Endpoint(), kOAuthToken)); + + auto provider = factory->CreateProvider(); + EXPECT_EQ(provider->GetAuthInfo(), kMockRootBuiltinToken); + + const int countBefore = Stub_.GetRequestCount(); + EXPECT_EQ(provider->GetAuthInfo(), kMockRootBuiltinToken); + EXPECT_EQ(Stub_.GetRequestCount(), countBefore); +} + +TEST(OAuth_WithFacility, ReturnsTokenFromMock) { + TIamTokenServiceStub stub; + stub.SetResponseToken(kMockRootBuiltinToken); + TIamGrpcServer server(&stub); + ASSERT_TRUE(server.Start()); + + auto factory = CreateIamOAuthCredentialsProviderFactory( + MakeOAuthParams(server.Endpoint(), kOAuthToken)); + auto facility = std::make_shared(); + auto provider = factory->CreateProvider(facility); + + EXPECT_EQ(provider->GetAuthInfo(), kMockRootBuiltinToken); + ASSERT_TRUE(stub.HasLastRequest()); + EXPECT_EQ(stub.GetLastRequest().yandex_passport_oauth_token(), kOAuthToken); +} diff --git a/tests/unit/client/CMakeLists.txt b/tests/unit/client/CMakeLists.txt index 7661422f540..b5574c44af6 100644 --- a/tests/unit/client/CMakeLists.txt +++ b/tests/unit/client/CMakeLists.txt @@ -58,6 +58,7 @@ add_ydb_test(NAME client-iam_grpc_ut GTEST LINK_LIBRARIES api-client-yc_public cpp-testing-common + tests-iam-mocks YDB-CPP-SDK::Iam LABELS unit @@ -70,6 +71,7 @@ add_ydb_test(NAME client-iam_http_ut GTEST http-server json cpp-testing-common + tests-iam-mocks YDB-CPP-SDK::Iam LABELS unit @@ -158,4 +160,89 @@ add_ydb_test(NAME client-ydb_spans_ut GTEST client-trace LABELS unit +) + +add_ydb_test(NAME client-build_info_ut GTEST + SOURCES + build_info/build_info_ut.cpp + LINK_LIBRARIES + YDB-CPP-SDK::Driver + LABELS + unit +) + +add_ydb_test(NAME client-query_session_ut + SOURCES + query/client_session_ut.cpp + LINK_LIBRARIES + client-ydb_query-impl + impl-session + library-operation_id + LABELS + unit +) + +add_ydb_test(NAME client-query_stats_ut + SOURCES + query/query_stats_ut.cpp + LINK_LIBRARIES + YDB-CPP-SDK::Query + LABELS + unit +) + +add_ydb_test(NAME client-row_ranges_ut + SOURCES + row_ranges/row_ranges_ut.cpp + LINK_LIBRARIES + YDB-CPP-SDK::Query + YDB-CPP-SDK::Table + LABELS + unit +) + +add_ydb_test(NAME client-row_ranges_public_include_ut + SOURCES + row_ranges/row_ranges_public_include_ut.cpp + LINK_LIBRARIES + YDB-CPP-SDK::Query + YDB-CPP-SDK::Table + LABELS + unit +) + +add_ydb_test(NAME client-retry_range_ut + SOURCES + retry_range/retry_range_ut.cpp + LINK_LIBRARIES + api-grpc + grpc-client + YDB-CPP-SDK::Driver + YDB-CPP-SDK::Query + YDB-CPP-SDK::Table + LABELS + unit +) + +add_ydb_test(NAME client-iam_private_grpc_ut GTEST + SOURCES + iam_private/grpc_iam_service_ut.cpp + LINK_LIBRARIES + api-client-yc_private + tests-iam-mocks + YDB-CPP-SDK::IamPrivate + LABELS + unit +) + +add_ydb_test(NAME client-metric_buffer_ut GTEST + INCLUDE_DIRS + ${YDB_SDK_SOURCE_DIR} + SOURCES + observability/metric_buffer_ut.cpp + LINK_LIBRARIES + impl-observability + client-metrics + LABELS + unit ) \ No newline at end of file diff --git a/tests/unit/client/endpoints/endpoints_ut.cpp b/tests/unit/client/endpoints/endpoints_ut.cpp index 7e2ebe719d4..69946265a73 100644 --- a/tests/unit/client/endpoints/endpoints_ut.cpp +++ b/tests/unit/client/endpoints/endpoints_ut.cpp @@ -7,6 +7,7 @@ #include #include +#include #include using namespace NYdb; @@ -144,6 +145,29 @@ Y_UNIT_TEST_SUITE(EndpointElector) { UNIT_ASSERT_VALUES_EQUAL(elector.GetPessimizationRatio(), 0); } + Y_UNIT_TEST(PessimizeNode) { + TEndpointElectorSafe elector; + elector.SetNewState(std::vector{ + {"host:2136", 1, "", 42}, + {"ipv4:10.0.0.1:2136", 1, "", 42}, + {"other:2136", 1, "", 99}, + }); + UNIT_ASSERT_VALUES_EQUAL(elector.GetPessimizationRatio(), 0); + + elector.PessimizeNode(42); + UNIT_ASSERT_VALUES_EQUAL(elector.GetPessimizationRatio(), 66); + + const auto maxPriority = std::numeric_limits::max(); + UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey("host:2136", 0), true).Priority, maxPriority); + UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey("ipv4:10.0.0.1:2136", 0), true).Priority, maxPriority); + UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey("", 42), true).Priority, maxPriority); + UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey("other:2136", 0), true).Priority, 1); + UNIT_ASSERT_VALUES_EQUAL(elector.GetEndpoint(TEndpointKey("", 99), true).Priority, 1); + + elector.PessimizeNode(42); + UNIT_ASSERT_VALUES_EQUAL(elector.GetPessimizationRatio(), 66); + } + Y_UNIT_TEST(Election) { TEndpointElectorSafe elector; elector.SetNewState(std::vector{{"Two", 2}, {"One_A", 1}, {"Three", 3}, {"One_B", 1}}); diff --git a/tests/unit/client/iam/grpc_iam_ut.cpp b/tests/unit/client/iam/grpc_iam_ut.cpp index a3c51f96194..e5eaa9ed1a9 100644 --- a/tests/unit/client/iam/grpc_iam_ut.cpp +++ b/tests/unit/client/iam/grpc_iam_ut.cpp @@ -1,13 +1,11 @@ #include #include #include +#include #include #include -#include -#include - #include #include @@ -19,118 +17,19 @@ #include #include +#include + using namespace NYdb; +using namespace NYdb::NTest; using namespace yandex::cloud::iam::v1; -namespace { - -class TBlockingIamTokenService final : public IamTokenService::Service { -public: - grpc::Status Create( - grpc::ServerContext* context, - const CreateIamTokenRequest*, - CreateIamTokenResponse* response) override - { - RpcEntered_.store(true); - - std::unique_lock lock(Mutex_); - ReleasedCv_.wait(lock, [this] { return Released_; }); - - if (context->IsCancelled()) { - return grpc::Status::CANCELLED; - } - response->set_iam_token("released-token"); - response->mutable_expires_at()->set_seconds(4102444800); - response->mutable_expires_at()->set_nanos(0); - return grpc::Status::OK; - } - - void Release() { - { - std::lock_guard lock(Mutex_); - Released_ = true; - } - ReleasedCv_.notify_all(); - } - - bool WaitUntilRpcEntered(std::chrono::milliseconds timeout) const { - const auto deadline = std::chrono::steady_clock::now() + timeout; - while (std::chrono::steady_clock::now() < deadline) { - if (RpcEntered_.load()) { - return true; - } - std::this_thread::sleep_for(std::chrono::milliseconds(2)); - } - return RpcEntered_.load(); - } - -private: - std::atomic RpcEntered_{false}; - bool Released_ = false; - mutable std::mutex Mutex_; - std::condition_variable ReleasedCv_; -}; - -class TIamGrpcServer { -public: - explicit TIamGrpcServer(TBlockingIamTokenService* service) - : Service_(service) - {} - - bool Start() { - grpc::ServerBuilder builder; - int boundPort = 0; - builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &boundPort); - builder.SetMaxSendMessageSize(4 * 1024 * 1024); - builder.RegisterService(Service_); - Server_ = builder.BuildAndStart(); - if (!Server_ || boundPort <= 0) { - return false; - } - Port_ = boundPort; - WaitThread_ = std::thread([this] { Server_->Wait(); }); - return true; - } - - std::string Endpoint() const { - return "127.0.0.1:" + std::to_string(Port_); - } - - void Stop() { - if (Server_) { - Server_->Shutdown(); - } - if (WaitThread_.joinable()) { - WaitThread_.join(); - } - Server_.reset(); - } - - ~TIamGrpcServer() { - Stop(); - } - -private: - TBlockingIamTokenService* Service_ = nullptr; - std::unique_ptr Server_; - int Port_ = 0; - std::thread WaitThread_; -}; - -} // namespace - TEST(GrpcIamCredentialsProvider, TeardownWhileIamCreatePendingCompletes) { TBlockingIamTokenService iamService; + TBlockingIamReleaseGuard releaseGuard(iamService); TIamGrpcServer server(&iamService); ASSERT_TRUE(server.Start()); - TIamOAuth params; - params.Endpoint = server.Endpoint(); - params.OAuthToken = "unit-test-oauth-token"; - params.EnableSsl = false; - params.RefreshPeriod = TDuration::Hours(1); - // Short client deadline: completion arrives with DEADLINE_EXCEEDED while the server handler is still - // blocked waiting for Release() (in-flight work on the server side vs client teardown). + TIamOAuth params = MakeOAuthParams(server.Endpoint()); params.RequestTimeout = TDuration::MilliSeconds(400); auto work = [¶ms] { @@ -157,14 +56,11 @@ TEST(GrpcIamCredentialsProvider, TeardownWhileIamCreatePendingCompletes) { TEST(GrpcIamCredentialsProvider, TeardownWhileIamCreatePendingCompletesViaFactoryWrapper) { TBlockingIamTokenService iamService; + TBlockingIamReleaseGuard releaseGuard(iamService); TIamGrpcServer server(&iamService); ASSERT_TRUE(server.Start()); - TIamOAuth params; - params.Endpoint = server.Endpoint(); - params.OAuthToken = "unit-test-oauth-token"; - params.EnableSsl = false; - params.RefreshPeriod = TDuration::Hours(1); + TIamOAuth params = MakeOAuthParams(server.Endpoint()); params.RequestTimeout = TDuration::MilliSeconds(400); auto factory = std::make_shared>(params); - - auto work = [&factory]() -> std::string { - auto provider = factory->CreateProvider(); - return provider->GetAuthInfo(); - }; - - std::future done = std::async(std::launch::async, work); - ASSERT_EQ(done.wait_for(std::chrono::seconds(20)), std::future_status::ready) - << "no-arg CreateProvider() path must produce a token and tear down cleanly"; - EXPECT_EQ(done.get(), "released-token"); - - server.Stop(); -} - -namespace { - -// PS256-compatible RSA test keypair (also used in jwt_token_source_ut.cpp). -constexpr const char* TestRSAPrivateKey = - "-----BEGIN PRIVATE KEY-----\n" - "MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQC75/JS3rMcLJxv\n" - "FgpOzF5+2gH+Yig3RE2MTl9uwC0BZKAv6foYr7xywQyWIK+W1cBhz8R4LfFmZo2j\n" - "M0aCvdRmNBdW0EDSTnHLxCsFhoQWLVq+bI5f5jzkcoiioUtaEpADPqwgVULVtN/n\n" - "nPJiZ6/dU30C3jmR6+LUgEntUtWt3eq3xQIn5lG3zC1klBY/HxtfH5Hu8xBvwRQT\n" - "Jnh3UpPLj8XwSmriDgdrhR7o6umWyVuGrMKlLHmeivlfzjYtfzO1MOIMG8t2/zxG\n" - "R+xb4Vwks73sH1KruH/0/JMXU97npwpe+Um+uXhpldPygGErEia7abyZB2gMpXqr\n" - "WYKMo02NAgMBAAECggEAO0BpC5OYw/4XN/optu4/r91bupTGHKNHlsIR2rDzoBhU\n" - "YLd1evpTQJY6O07EP5pYZx9mUwUdtU4KRJeDGO/1/WJYp7HUdtxwirHpZP0lQn77\n" - "uccuX/QQaHLrPekBgz4ONk+5ZBqukAfQgM7fKYOLk41jgpeDbM2Ggb6QUSsJISEp\n" - "zrwpI/nNT/wn+Hvx4DxrzWU6wF+P8kl77UwPYlTA7GsT+T7eKGVH8xsxmK8pt6lg\n" - "svlBA5XosWBWUCGLgcBkAY5e4ZWbkdd183o+oMo78id6C+PQPE66PLDtHWfpRRmN\n" - "m6XC03x6NVhnfvfozoWnmS4+e4qj4F/emCHvn0GMywKBgQDLXlj7YPFVXxZpUvg/\n" - "rheVcCTGbNmQJ+4cZXx87huqwqKgkmtOyeWsRc7zYInYgraDrtCuDBCfP//ZzOh0\n" - "LxepYLTPk5eNn/GT+VVrqsy35Ccr60g7Lp/bzb1WxyhcLbo0KX7/6jl0lP+VKtdv\n" - "mto+4mbSBXSM1Y5BVVoVgJ3T/wKBgQDsiSvPRzVi5TTj13x67PFymTMx3HCe2WzH\n" - "JUyepCmVhTm482zW95pv6raDr5CTO6OYpHtc5sTTRhVYEZoEYFTM9Vw8faBtluWG\n" - "BjkRh4cIpoIARMn74YZKj0C/0vdX7SHdyBOU3bgRPHg08Hwu3xReqT1kEPSI/B2V\n" - "4pe5fVrucwKBgQCNFgUxUA3dJjyMES18MDDYUZaRug4tfiYouRdmLGIxUxozv6CG\n" - "ZnbZzwxFt+GpvPUV4f+P33rgoCvFU+yoPctyjE6j+0aW0DFucPmb2kBwCu5J/856\n" - "kFwCx3blbwFHAco+SdN7g2kcwgmV2MTg/lMOcU7XwUUcN0Obe7UlWbckzQKBgQDQ\n" - "nXaXHL24GGFaZe4y2JFmujmNy1dEsoye44W9ERpf9h1fwsoGmmCKPp90az5+rIXw\n" - "FXl8CUgk8lXW08db/r4r+ma8Lyx0GzcZyplAnaB5/6j+pazjSxfO4KOBy4Y89Tb+\n" - "TP0AOcCi6ws13bgY+sUTa/5qKA4UVw+c5zlb7nRpgwKBgGXAXhenFw1666482iiN\n" - "cHSgwc4ZHa1oL6aNJR1XWH+aboBSwR+feKHUPeT4jHgzRGo/aCNHD2FE5I8eBv33\n" - "of1kWYjAO0YdzeKrW0rTwfvt9gGg+CS397aWu4cy+mTI+MNfBgeDAIVBeJOJXLlX\n" - "hL8bFAuNNVrCOp79TNnNIsh7\n" - "-----END PRIVATE KEY-----\n"; - -constexpr const char* TestRSAPublicKey = - "-----BEGIN PUBLIC KEY-----\n" - "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAu+fyUt6zHCycbxYKTsxe\n" - "ftoB/mIoN0RNjE5fbsAtAWSgL+n6GK+8csEMliCvltXAYc/EeC3xZmaNozNGgr3U\n" - "ZjQXVtBA0k5xy8QrBYaEFi1avmyOX+Y85HKIoqFLWhKQAz6sIFVC1bTf55zyYmev\n" - "3VN9At45kevi1IBJ7VLVrd3qt8UCJ+ZRt8wtZJQWPx8bXx+R7vMQb8EUEyZ4d1KT\n" - "y4/F8Epq4g4Ha4Ue6OrplslbhqzCpSx5nor5X842LX8ztTDiDBvLdv88RkfsW+Fc\n" - "JLO97B9Sq7h/9PyTF1Pe56cKXvlJvrl4aZXT8oBhKxImu2m8mQdoDKV6q1mCjKNN\n" - "jQIDAQAB\n" - "-----END PUBLIC KEY-----\n"; - -} // namespace - -// Regression test for the deprecated no-arg CreateProvider() on the JWT factory. Mirrors the -// OAuth counterpart above. The standalone path spins up a private TSimpleCoreFacility behind -// TOwningFacilityCredentialsProvider; this verifies the JWT specialization wires it up -// correctly and tears down without aborting. -TEST(GrpcIamCredentialsProvider, NoArgCreateProviderBackwardCompatJwt) { - TBlockingIamTokenService iamService; - TIamGrpcServer server(&iamService); - ASSERT_TRUE(server.Start()); - - iamService.Release(); - - TIamJwtParams params; - params.Endpoint = server.Endpoint(); - params.EnableSsl = false; - params.RefreshPeriod = TDuration::Hours(1); - params.RequestTimeout = TDuration::Seconds(5); - params.JwtParams.AccountId = "unit-test-account"; - params.JwtParams.KeyId = "unit-test-key"; - params.JwtParams.PrivKey = TestRSAPrivateKey; - params.JwtParams.PubKey = TestRSAPublicKey; - - auto factory = std::make_shared>(params); - - auto work = [&factory]() -> std::string { - auto provider = factory->CreateProvider(); - return provider->GetAuthInfo(); - }; - - std::future done = std::async(std::launch::async, work); - ASSERT_EQ(done.wait_for(std::chrono::seconds(20)), std::future_status::ready) - << "no-arg CreateProvider() path must produce a token and tear down cleanly"; - EXPECT_EQ(done.get(), "released-token"); - - server.Stop(); -} - namespace { class TSlowBlockingAuthProvider final : public ICredentialsProvider { @@ -350,24 +128,45 @@ class TSlowBlockingAuthProvider final : public ICredentialsProvider { bool Released_ = false; }; +class TFailThenSucceedAuthProvider final : public ICredentialsProvider { +public: + explicit TFailThenSucceedAuthProvider(int failCount) + : FailCount_(failCount) + {} + + std::string GetAuthInfo() const override { + if (CallCount_.fetch_add(1) < FailCount_) { + ythrow yexception() << "auth failure"; + } + return "auth-token"; + } + + bool IsValid() const override { + return true; + } + + int GetCallCount() const { + return CallCount_.load(); + } + +private: + const int FailCount_; + mutable std::atomic CallCount_{0}; +}; + } // namespace -// Regression: Stop() must not hang when FillContext is blocked inside AuthTokenProvider_->GetAuthInfo(). TEST(GrpcIamCredentialsProvider, StopDuringFillContextDoesNotHang) { - TBlockingIamTokenService iamService; - TIamGrpcServer server(&iamService); + TIamTokenServiceStub iamStub; + iamStub.SetResponseToken("unit-test-iam-token"); + TIamGrpcServer server(&iamStub); ASSERT_TRUE(server.Start()); - iamService.Release(); - auto authProvider = std::make_shared(); - TIamOAuth params; - params.Endpoint = server.Endpoint(); - params.OAuthToken = "unit-test-oauth-token"; - params.EnableSsl = false; + TIamOAuth params = MakeOAuthParams(server.Endpoint()); params.RefreshPeriod = TDuration::MilliSeconds(50); - params.RequestTimeout = TDuration::Seconds(5); + params.RequestTimeout = TDuration::MilliSeconds(400); std::shared_ptr> provider; auto facility = std::make_shared(); @@ -384,21 +183,20 @@ TEST(GrpcIamCredentialsProvider, StopDuringFillContextDoesNotHang) { facility, authProvider); - ASSERT_TRUE(authProvider->WaitUntilBlocked(std::chrono::seconds(10))) + struct TAuthReleaseGuard { + std::shared_ptr Provider; + ~TAuthReleaseGuard() { + if (Provider) { + Provider->Release(); + } + } + } authReleaseGuard{authProvider}; + + ASSERT_TRUE(authProvider->WaitUntilBlocked(std::chrono::seconds(30))) << "FillContext should block inside slow AuthTokenProvider during refresh"; - // Move-capture provider so provider.reset() is the last owner and triggers Stop(). - // facility stays here: its destructor joins the worker thread (still blocked in - // GetAuthInfo()), so it must be destroyed after authProvider->Release(). - // - // Mirrors production (TOwningFacilityCredentialsProvider): Inner_ dies first (Stop()), - // then Facility_ dies (joins worker). No self-join: TImpl holds facility as weak_ptr, - // callbacks capture only weak_ptr. If a callback locks TImpl temporarily, TImpl's - // destructor is trivial (no joins), and it never strong-refs the facility. std::future stopDone = std::async(std::launch::async, [provider = std::move(provider)]() mutable { - // Last owner: triggers ~TGrpcIamCredentialsProvider -> Stop(). - // Stop() must return even though the worker thread is still blocked in GetAuthInfo(). provider.reset(); }); @@ -406,9 +204,35 @@ TEST(GrpcIamCredentialsProvider, StopDuringFillContextDoesNotHang) { << "provider destructor (Stop()) must complete while FillContext is blocked in GetAuthInfo()"; stopDone.get(); - // Now unblock the worker thread so facility can be destroyed cleanly. - authProvider->Release(); - facility.reset(); // joins the worker thread (now unblocked) + server.Stop(); +} + +TEST(GrpcIamCredentialsProvider, FillContextAuthExceptionSurvivesAndRecovers) { + TIamTokenServiceStub iamStub; + iamStub.SetResponseToken("unit-test-iam-token"); + TIamGrpcServer server(&iamStub); + ASSERT_TRUE(server.Start()); + + auto authProvider = std::make_shared(2); + + TIamOAuth params = MakeOAuthParams(server.Endpoint()); + auto facility = std::make_shared(); + + TGrpcIamCredentialsProvider provider( + params, + [token = params.OAuthToken](CreateIamTokenRequest& req) { + req.set_yandex_passport_oauth_token(TStringType{token}); + }, + [](IamTokenService::Stub* stub, grpc::ClientContext* context, const CreateIamTokenRequest* request, + CreateIamTokenResponse* response, std::function cb) { + stub->async()->Create(context, request, response, std::move(cb)); + }, + facility, + authProvider); + + EXPECT_EQ(provider.GetAuthInfo(), "unit-test-iam-token"); + EXPECT_EQ(iamStub.GetRequestCount(), 1); + EXPECT_GE(authProvider->GetCallCount(), 3); server.Stop(); } diff --git a/tests/unit/client/iam/http_iam_ut.cpp b/tests/unit/client/iam/http_iam_ut.cpp index 630d87618fd..f2ed793fe85 100644 --- a/tests/unit/client/iam/http_iam_ut.cpp +++ b/tests/unit/client/iam/http_iam_ut.cpp @@ -1,160 +1,29 @@ #include -#include -#include -#include - -#include +#include #include #include -#include -#include #include #include using namespace NYdb; - -class TMetadataServer : public THttpServer::ICallBack { -public: - class TRequest : public TRequestReplier { - public: - explicit TRequest(TMetadataServer* server) - : Server(server) - {} - - bool DoReply(const TReplyParams& params) override { - { - std::lock_guard lock(Server->Lock); - ++Server->RequestCount; - } - - THttpResponse resp(Server->StatusCode); - resp.SetContent(TString{Server->Response}); - resp.OutTo(params.Output); - return true; - } - - TMetadataServer* Server = nullptr; - }; - - TMetadataServer() - : PortHolder(NTesting::GetFreePort()) - , Port(static_cast(PortHolder)) - , HttpOptions(Port) - , HttpServer(this, HttpOptions) - { - HttpServer.Start(); - } - - ~TMetadataServer() { - HttpServer.Stop(); - } - - TClientRequest* CreateClient() override { - return new TRequest(this); - } - - void SetResponse(HttpCodes code, const std::string& response) { - StatusCode = code; - Response = response; - } - - int GetRequestCount() const { - std::lock_guard lock(Lock); - return RequestCount; - } - - void ResetRequestCount() { - std::lock_guard lock(Lock); - RequestCount = 0; - } - - NTesting::TPortHolder PortHolder; - uint16_t Port; - THttpServer::TOptions HttpOptions; - THttpServer HttpServer; - HttpCodes StatusCode = HTTP_OK; - std::string Response; - mutable std::mutex Lock; - int RequestCount = 0; -}; - -static std::string MakeTokenResponse(const std::string& token, int expiresIn) { - TStringStream ss; - NJson::TJsonWriter w(&ss, false); - w.OpenMap(); - w.WriteKey("access_token"); - w.Write(token); - w.WriteKey("token_type"); - w.Write("Bearer"); - w.WriteKey("expires_in"); - w.Write(expiresIn); - w.CloseMap(); - w.Flush(); - return ss.Str(); -} - -static std::string MakeTokenResponseWithExpiry(const std::string& token, const TInstant& expiry) { - TStringStream ss; - NJson::TJsonWriter w(&ss, false); - w.OpenMap(); - w.WriteKey("access_token"); - w.Write(token); - w.WriteKey("token_type"); - w.Write("Bearer"); - w.WriteKey("expiry"); - w.Write(expiry.FormatGmTime("%Y-%m-%dT%H:%M:%S.000000000Z")); - w.CloseMap(); - w.Flush(); - return ss.Str(); -} - -static std::string MakeTokenResponseNoExpiry(const std::string& token) { - TStringStream ss; - NJson::TJsonWriter w(&ss, false); - w.OpenMap(); - w.WriteKey("access_token"); - w.Write(token); - w.WriteKey("token_type"); - w.Write("Bearer"); - w.CloseMap(); - w.Flush(); - return ss.Str(); -} - -TEST(IamCredentialsProvider, BasicTokenFetch) { - TMetadataServer server; - server.SetResponse(HTTP_OK, MakeTokenResponse("test-token-123", 3600)); - - TIamHost params; - params.Host = "localhost"; - params.Port = server.Port; - params.RefreshPeriod = TDuration::Hours(1); - - auto factory = CreateIamCredentialsProviderFactory(params); - auto provider = factory->CreateProvider(); - - EXPECT_EQ(provider->GetAuthInfo(), "test-token-123"); -} +using namespace NYdb::NTest; TEST(IamCredentialsProvider, ExpiryFieldSupport) { TMetadataServer server; + server.SetStrictMode(false); auto expiry = TInstant::Now() + TDuration::Hours(12); server.SetResponse(HTTP_OK, MakeTokenResponseWithExpiry("expiry-token", expiry)); - TIamHost params; - params.Host = "localhost"; - params.Port = server.Port; - params.RefreshPeriod = TDuration::Hours(1); + TIamHost params = MakeMetadataParams(server.Port); auto factory = CreateIamCredentialsProviderFactory(params); auto provider = factory->CreateProvider(); EXPECT_EQ(provider->GetAuthInfo(), "expiry-token"); - // Second call should not trigger refresh (token is still valid) int countBefore = server.GetRequestCount(); provider->GetAuthInfo(); EXPECT_EQ(server.GetRequestCount(), countBefore); @@ -162,20 +31,16 @@ TEST(IamCredentialsProvider, ExpiryFieldSupport) { TEST(IamCredentialsProvider, NoExpiryFieldFallback) { TMetadataServer server; + server.SetStrictMode(false); server.SetResponse(HTTP_OK, MakeTokenResponseNoExpiry("no-expiry-token")); - TIamHost params; - params.Host = "localhost"; - params.Port = server.Port; - params.RefreshPeriod = TDuration::Hours(1); + TIamHost params = MakeMetadataParams(server.Port); auto factory = CreateIamCredentialsProviderFactory(params); auto provider = factory->CreateProvider(); - // Token should be saved even without expires_in/expiry EXPECT_EQ(provider->GetAuthInfo(), "no-expiry-token"); - // Should not immediately refresh (fallback interval should be > 0) int countBefore = server.GetRequestCount(); provider->GetAuthInfo(); EXPECT_EQ(server.GetRequestCount(), countBefore); @@ -183,27 +48,24 @@ TEST(IamCredentialsProvider, NoExpiryFieldFallback) { TEST(IamCredentialsProvider, ServerError) { TMetadataServer server; + server.SetStrictMode(false); server.SetResponse(HTTP_INTERNAL_SERVER_ERROR, ""); - TIamHost params; - params.Host = "localhost"; - params.Port = server.Port; + TIamHost params = MakeMetadataParams(server.Port); auto factory = CreateIamCredentialsProviderFactory(params); auto provider = factory->CreateProvider(); - // Constructor GetTicket() fails, token should be empty EXPECT_EQ(provider->GetAuthInfo(), ""); } TEST(IamCredentialsProvider, ConcurrentAccess) { TMetadataServer server; + server.SetStrictMode(false); server.SetResponse(HTTP_OK, MakeTokenResponse("concurrent-token", 3600)); - TIamHost params; - params.Host = "localhost"; - params.Port = server.Port; - params.RefreshPeriod = TDuration::MilliSeconds(1); // Force frequent refreshes + TIamHost params = MakeMetadataParams(server.Port); + params.RefreshPeriod = TDuration::MilliSeconds(1); auto factory = CreateIamCredentialsProviderFactory(params); auto provider = factory->CreateProvider(); @@ -211,11 +73,12 @@ TEST(IamCredentialsProvider, ConcurrentAccess) { constexpr int NUM_THREADS = 8; constexpr int ITERATIONS = 100; - std::vector> threads; + std::vector threads; + threads.reserve(NUM_THREADS); std::atomic errors{0}; for (int i = 0; i < NUM_THREADS; ++i) { - threads.push_back(std::make_unique([&]() { + threads.emplace_back([&]() { for (int j = 0; j < ITERATIONS; ++j) { try { auto token = provider->GetAuthInfo(); @@ -226,11 +89,11 @@ TEST(IamCredentialsProvider, ConcurrentAccess) { errors.fetch_add(1); } } - })); + }); } - for (auto& t : threads) { - t->join(); + for (auto& thread : threads) { + thread.join(); } EXPECT_EQ(errors.load(), 0); diff --git a/tests/unit/client/iam_private/grpc_iam_service_ut.cpp b/tests/unit/client/iam_private/grpc_iam_service_ut.cpp index 1478a9997d1..f557bd804c5 100644 --- a/tests/unit/client/iam_private/grpc_iam_service_ut.cpp +++ b/tests/unit/client/iam_private/grpc_iam_service_ut.cpp @@ -1,12 +1,10 @@ #include #include +#include #include #include -#include -#include - #include #include @@ -16,51 +14,11 @@ #include using namespace NYdb; +using namespace NYdb::NTest; using namespace yandex::cloud::priv::iam::v1; namespace { -constexpr const char* TestRSAPrivateKey = - "-----BEGIN PRIVATE KEY-----\n" - "MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQC75/JS3rMcLJxv\n" - "FgpOzF5+2gH+Yig3RE2MTl9uwC0BZKAv6foYr7xywQyWIK+W1cBhz8R4LfFmZo2j\n" - "M0aCvdRmNBdW0EDSTnHLxCsFhoQWLVq+bI5f5jzkcoiioUtaEpADPqwgVULVtN/n\n" - "nPJiZ6/dU30C3jmR6+LUgEntUtWt3eq3xQIn5lG3zC1klBY/HxtfH5Hu8xBvwRQT\n" - "Jnh3UpPLj8XwSmriDgdrhR7o6umWyVuGrMKlLHmeivlfzjYtfzO1MOIMG8t2/zxG\n" - "R+xb4Vwks73sH1KruH/0/JMXU97npwpe+Um+uXhpldPygGErEia7abyZB2gMpXqr\n" - "WYKMo02NAgMBAAECggEAO0BpC5OYw/4XN/optu4/r91bupTGHKNHlsIR2rDzoBhU\n" - "YLd1evpTQJY6O07EP5pYZx9mUwUdtU4KRJeDGO/1/WJYp7HUdtxwirHpZP0lQn77\n" - "uccuX/QQaHLrPekBgz4ONk+5ZBqukAfQgM7fKYOLk41jgpeDbM2Ggb6QUSsJISEp\n" - "zrwpI/nNT/wn+Hvx4DxrzWU6wF+P8kl77UwPYlTA7GsT+T7eKGVH8xsxmK8pt6lg\n" - "svlBA5XosWBWUCGLgcBkAY5e4ZWbkdd183o+oMo78id6C+PQPE66PLDtHWfpRRmN\n" - "m6XC03x6NVhnfvfozoWnmS4+e4qj4F/emCHvn0GMywKBgQDLXlj7YPFVXxZpUvg/\n" - "rheVcCTGbNmQJ+4cZXx87huqwqKgkmtOyeWsRc7zYInYgraDrtCuDBCfP//ZzOh0\n" - "LxepYLTPk5eNn/GT+VVrqsy35Ccr60g7Lp/bzb1WxyhcLbo0KX7/6jl0lP+VKtdv\n" - "mto+4mbSBXSM1Y5BVVoVgJ3T/wKBgQDsiSvPRzVi5TTj13x67PFymTMx3HCe2WzH\n" - "JUyepCmVhTm482zW95pv6raDr5CTO6OYpHtc5sTTRhVYEZoEYFTM9Vw8faBtluWG\n" - "BjkRh4cIpoIARMn74YZKj0C/0vdX7SHdyBOU3bgRPHg08Hwu3xReqT1kEPSI/B2V\n" - "4pe5fVrucwKBgQCNFgUxUA3dJjyMES18MDDYUZaRug4tfiYouRdmLGIxUxozv6CG\n" - "ZnbZzwxFt+GpvPUV4f+P33rgoCvFU+yoPctyjE6j+0aW0DFucPmb2kBwCu5J/856\n" - "kFwCx3blbwFHAco+SdN7g2kcwgmV2MTg/lMOcU7XwUUcN0Obe7UlWbckzQKBgQDQ\n" - "nXaXHL24GGFaZe4y2JFmujmNy1dEsoye44W9ERpf9h1fwsoGmmCKPp90az5+rIXw\n" - "FXl8CUgk8lXW08db/r4r+ma8Lyx0GzcZyplAnaB5/6j+pazjSxfO4KOBy4Y89Tb+\n" - "TP0AOcCi6ws13bgY+sUTa/5qKA4UVw+c5zlb7nRpgwKBgGXAXhenFw1666482iiN\n" - "cHSgwc4ZHa1oL6aNJR1XWH+aboBSwR+feKHUPeT4jHgzRGo/aCNHD2FE5I8eBv33\n" - "of1kWYjAO0YdzeKrW0rTwfvt9gGg+CS397aWu4cy+mTI+MNfBgeDAIVBeJOJXLlX\n" - "hL8bFAuNNVrCOp79TNnNIsh7\n" - "-----END PRIVATE KEY-----\n"; - -constexpr const char* TestRSAPublicKey = - "-----BEGIN PUBLIC KEY-----\n" - "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAu+fyUt6zHCycbxYKTsxe\n" - "ftoB/mIoN0RNjE5fbsAtAWSgL+n6GK+8csEMliCvltXAYc/EeC3xZmaNozNGgr3U\n" - "ZjQXVtBA0k5xy8QrBYaEFi1avmyOX+Y85HKIoqFLWhKQAz6sIFVC1bTf55zyYmev\n" - "3VN9At45kevi1IBJ7VLVrd3qt8UCJ+ZRt8wtZJQWPx8bXx+R7vMQb8EUEyZ4d1KT\n" - "y4/F8Epq4g4Ha4Ue6OrplslbhqzCpSx5nor5X842LX8ztTDiDBvLdv88RkfsW+Fc\n" - "JLO97B9Sq7h/9PyTF1Pe56cKXvlJvrl4aZXT8oBhKxImu2m8mQdoDKV6q1mCjKNN\n" - "jQIDAQAB\n" - "-----END PUBLIC KEY-----\n"; - // Implements both Create (used by the nested JWT provider) and CreateForService (used by the // outer service provider). Returns distinct tokens so the outer one is observable via GetAuthInfo(). class TIamServiceStub final : public IamTokenService::Service { @@ -86,45 +44,6 @@ class TIamServiceStub final : public IamTokenService::Service { } }; -class TIamGrpcServer { -public: - explicit TIamGrpcServer(TIamServiceStub* service) : Service_(service) {} - - bool Start() { - grpc::ServerBuilder builder; - int boundPort = 0; - builder.AddListeningPort("127.0.0.1:0", grpc::InsecureServerCredentials(), &boundPort); - builder.RegisterService(Service_); - Server_ = builder.BuildAndStart(); - if (!Server_ || boundPort <= 0) { - return false; - } - Port_ = boundPort; - WaitThread_ = std::thread([this] { Server_->Wait(); }); - return true; - } - - std::string Endpoint() const { return "127.0.0.1:" + std::to_string(Port_); } - - void Stop() { - if (Server_) { - Server_->Shutdown(); - } - if (WaitThread_.joinable()) { - WaitThread_.join(); - } - Server_.reset(); - } - - ~TIamGrpcServer() { Stop(); } - -private: - TIamServiceStub* Service_ = nullptr; - std::unique_ptr Server_; - int Port_ = 0; - std::thread WaitThread_; -}; - } // namespace // Regression test for the deprecated no-arg CreateProvider() on the IAM service-account @@ -138,16 +57,7 @@ TEST(IamServiceCredentialsProvider, NoArgCreateProviderWithGrpcInnerCreds) { TIamGrpcServer server(&stub); ASSERT_TRUE(server.Start()); - TIamJwtParams jwtParams; - jwtParams.Endpoint = server.Endpoint(); - jwtParams.EnableSsl = false; - jwtParams.RefreshPeriod = TDuration::Hours(1); - jwtParams.RequestTimeout = TDuration::Seconds(5); - jwtParams.JwtParams.AccountId = "unit-test-account"; - jwtParams.JwtParams.KeyId = "unit-test-key"; - jwtParams.JwtParams.PrivKey = TestRSAPrivateKey; - jwtParams.JwtParams.PubKey = TestRSAPublicKey; - + auto jwtParams = MakeJwtParams(server.Endpoint()); auto jwtFactory = std::make_shared>(jwtParams); @@ -185,16 +95,7 @@ TEST(IamServiceCredentialsProvider, NoArgCreateProviderTemporaryFactory) { TIamGrpcServer server(&stub); ASSERT_TRUE(server.Start()); - TIamJwtParams jwtParams; - jwtParams.Endpoint = server.Endpoint(); - jwtParams.EnableSsl = false; - jwtParams.RefreshPeriod = TDuration::Hours(1); - jwtParams.RequestTimeout = TDuration::Seconds(5); - jwtParams.JwtParams.AccountId = "unit-test-account"; - jwtParams.JwtParams.KeyId = "unit-test-key"; - jwtParams.JwtParams.PrivKey = TestRSAPrivateKey; - jwtParams.JwtParams.PubKey = TestRSAPublicKey; - + auto jwtParams = MakeJwtParams(server.Endpoint()); auto jwtFactory = std::make_shared>(jwtParams); diff --git a/tests/unit/client/query/client_session_ut.cpp b/tests/unit/client/query/client_session_ut.cpp new file mode 100644 index 00000000000..cf9610f2390 --- /dev/null +++ b/tests/unit/client/query/client_session_ut.cpp @@ -0,0 +1,136 @@ +#include + +#include + +#include +#include + +using namespace NYdb; +using namespace NYdb::NQuery; + +namespace { + +class TMockSessionClient : public ISessionClient { +public: + void DeleteSession(TKqpSessionCommon*) override { + } + + void PessimizeNode(std::uint64_t nodeId) override { + PessimizedNodeId = nodeId; + ++PessimizeCalls; + } + + bool ReturnSession(TKqpSessionCommon*) override { + return true; + } + + std::uint64_t PessimizedNodeId = 0; + int PessimizeCalls = 0; +}; + +class TMockServerCloseHandler : public IServerCloseHandler { +public: + void OnCloseSession(const TKqpSessionCommon*, std::shared_ptr) override { + ++CloseCalls; + } + + int CloseCalls = 0; +}; + +std::string MakeSessionIdWithNodeId(std::uint64_t nodeId) { + NKikimr::NOperationId::TOperationId operationId; + operationId.SetKind(NKikimr::NOperationId::TOperationId::SESSION_YQL); + operationId.AddOptionalValue("node_id", ToString(nodeId)); + return operationId.ToString(); +} + +class TTestKqpSession : public TKqpSessionCommon { +public: + TTestKqpSession(const std::string& sessionId, const std::string& endpoint) + : TKqpSessionCommon(sessionId, endpoint, true) + { + MarkActive(); + } +}; + +Ydb::Query::SessionState MakeSessionShutdownState() { + Ydb::Query::SessionState state; + state.mutable_session_shutdown(); + return state; +} + +Ydb::Query::SessionState MakeNodeShutdownState() { + Ydb::Query::SessionState state; + state.mutable_node_shutdown(); + return state; +} + +} // namespace + +Y_UNIT_TEST_SUITE(QueryAttachSessionState) { + +Y_UNIT_TEST(SessionShutdownActiveSessionMarksClosing) { + TTestKqpSession session(MakeSessionIdWithNodeId(42), "host:2136"); + auto client = std::make_shared(); + + UNIT_ASSERT(HandleAttachSessionState(MakeSessionShutdownState(), &session, client) + == EAttachStreamReadAction::Stop); + UNIT_ASSERT(session.GetState() == TKqpSessionCommon::S_CLOSING); + UNIT_ASSERT_VALUES_EQUAL(client->PessimizeCalls, 0); +} + +Y_UNIT_TEST(SessionShutdownIdleInPoolDelegatesToCloseHandler) { + TTestKqpSession session(MakeSessionIdWithNodeId(42), "host:2136"); + auto client = std::make_shared(); + TMockServerCloseHandler closeHandler; + session.MarkIdle(); + session.UpdateServerCloseHandler(&closeHandler); + + UNIT_ASSERT(HandleAttachSessionState(MakeSessionShutdownState(), &session, client) + == EAttachStreamReadAction::Stop); + UNIT_ASSERT_VALUES_EQUAL(closeHandler.CloseCalls, 1); + UNIT_ASSERT_VALUES_EQUAL(client->PessimizeCalls, 0); +} + +Y_UNIT_TEST(NodeShutdownActiveSessionMarksClosingAndPessimizesNode) { + TTestKqpSession session(MakeSessionIdWithNodeId(42), "host:2136"); + auto client = std::make_shared(); + + UNIT_ASSERT(HandleAttachSessionState(MakeNodeShutdownState(), &session, client) + == EAttachStreamReadAction::Stop); + UNIT_ASSERT(session.GetState() == TKqpSessionCommon::S_CLOSING); + UNIT_ASSERT_VALUES_EQUAL(client->PessimizeCalls, 1); + UNIT_ASSERT_VALUES_EQUAL(client->PessimizedNodeId, 42U); +} + +Y_UNIT_TEST(NodeShutdownWithZeroNodeIdSkipsPessimization) { + TTestKqpSession session("", "host:2136"); + auto client = std::make_shared(); + + UNIT_ASSERT(session.GetEndpointKey().GetNodeId() == 0U); + UNIT_ASSERT(HandleAttachSessionState(MakeNodeShutdownState(), &session, client) + == EAttachStreamReadAction::Stop); + UNIT_ASSERT(session.GetState() == TKqpSessionCommon::S_CLOSING); + UNIT_ASSERT_VALUES_EQUAL(client->PessimizeCalls, 0); +} + +Y_UNIT_TEST(EmptySessionStateContinuesReading) { + TTestKqpSession session(MakeSessionIdWithNodeId(7), "host:2136"); + auto client = std::make_shared(); + Ydb::Query::SessionState state; + + UNIT_ASSERT(HandleAttachSessionState(state, &session, client) + == EAttachStreamReadAction::Continue); + UNIT_ASSERT(session.GetState() == TKqpSessionCommon::S_ACTIVE); + UNIT_ASSERT_VALUES_EQUAL(client->PessimizeCalls, 0); +} + +Y_UNIT_TEST(SessionShutdownNullSessionStopsReading) { + auto client = std::make_shared(); + + UNIT_ASSERT(HandleAttachSessionState(MakeSessionShutdownState(), nullptr, client) + == EAttachStreamReadAction::Stop); + UNIT_ASSERT_VALUES_EQUAL(client->PessimizeCalls, 0); +} + +} diff --git a/tests/unit/client/query/query_stats_ut.cpp b/tests/unit/client/query/query_stats_ut.cpp new file mode 100644 index 00000000000..9efb1c1a3aa --- /dev/null +++ b/tests/unit/client/query/query_stats_ut.cpp @@ -0,0 +1,91 @@ +#include + +#include + +#include + +using namespace NYdb::NQuery; + +Y_UNIT_TEST_SUITE(QueryStats) { + Y_UNIT_TEST(MapsProtoValues) { + Ydb::TableStats::QueryStats proto; + proto.set_process_cpu_time_us(10); + proto.set_total_duration_us(20); + proto.set_total_cpu_time_us(30); + proto.set_query_plan("plan"); + proto.set_query_ast("ast"); + proto.set_query_meta("meta"); + + auto* compilation = proto.mutable_compilation(); + compilation->set_from_cache(true); + compilation->set_duration_us(40); + compilation->set_cpu_time_us(50); + + auto* phase = proto.add_query_phases(); + phase->set_duration_us(60); + phase->set_cpu_time_us(70); + phase->set_affected_shards(80); + phase->set_literal_phase(true); + + auto* table = phase->add_table_access(); + table->set_name("table"); + table->set_partitions_count(90); + table->mutable_reads()->set_rows(100); + table->mutable_reads()->set_bytes(110); + table->mutable_updates()->set_rows(120); + table->mutable_updates()->set_bytes(130); + table->mutable_deletes()->set_rows(140); + table->mutable_deletes()->set_bytes(150); + + const TExecStats stats(proto); + UNIT_ASSERT_VALUES_EQUAL(stats.GetProcessCpuTimeUs(), 10); + UNIT_ASSERT_VALUES_EQUAL(stats.GetProcessCpuTime().MicroSeconds(), 10); + UNIT_ASSERT_VALUES_EQUAL(stats.GetTotalDurationUs(), 20); + UNIT_ASSERT_VALUES_EQUAL(stats.GetTotalDuration().MicroSeconds(), 20); + UNIT_ASSERT_VALUES_EQUAL(stats.GetTotalCpuTimeUs(), 30); + UNIT_ASSERT_VALUES_EQUAL(stats.GetTotalCpuTime().MicroSeconds(), 30); + const auto plan = stats.GetPlan(); + const auto ast = stats.GetAst(); + const auto meta = stats.GetMeta(); + UNIT_ASSERT_VALUES_EQUAL(*plan, "plan"); + UNIT_ASSERT_VALUES_EQUAL(*ast, "ast"); + UNIT_ASSERT_VALUES_EQUAL(*meta, "meta"); + + const auto compilationStats = stats.GetCompilation(); + UNIT_ASSERT(compilationStats); + UNIT_ASSERT(compilationStats->IsFromCache()); + UNIT_ASSERT_VALUES_EQUAL(compilationStats->GetDurationUs(), 40); + UNIT_ASSERT_VALUES_EQUAL(compilationStats->GetDuration().MicroSeconds(), 40); + UNIT_ASSERT_VALUES_EQUAL(compilationStats->GetCpuTimeUs(), 50); + UNIT_ASSERT_VALUES_EQUAL(compilationStats->GetCpuTime().MicroSeconds(), 50); + + const auto phases = stats.GetQueryPhases(); + UNIT_ASSERT_VALUES_EQUAL(phases.size(), 1); + UNIT_ASSERT_VALUES_EQUAL(phases[0].GetDurationUs(), 60); + UNIT_ASSERT_VALUES_EQUAL(phases[0].GetDuration().MicroSeconds(), 60); + UNIT_ASSERT_VALUES_EQUAL(phases[0].GetCpuTimeUs(), 70); + UNIT_ASSERT_VALUES_EQUAL(phases[0].GetCpuTime().MicroSeconds(), 70); + UNIT_ASSERT_VALUES_EQUAL(phases[0].GetAffectedShards(), 80); + UNIT_ASSERT(phases[0].IsLiteralPhase()); + + const auto& tables = phases[0].GetTableAccess(); + UNIT_ASSERT_VALUES_EQUAL(tables.size(), 1); + UNIT_ASSERT_VALUES_EQUAL(tables[0].GetName(), "table"); + UNIT_ASSERT_VALUES_EQUAL(tables[0].GetPartitionsCount(), 90); + UNIT_ASSERT_VALUES_EQUAL(tables[0].GetReads().GetRows(), 100); + UNIT_ASSERT_VALUES_EQUAL(tables[0].GetReads().GetBytes(), 110); + UNIT_ASSERT_VALUES_EQUAL(tables[0].GetUpdates().GetRows(), 120); + UNIT_ASSERT_VALUES_EQUAL(tables[0].GetUpdates().GetBytes(), 130); + UNIT_ASSERT_VALUES_EQUAL(tables[0].GetDeletes().GetRows(), 140); + UNIT_ASSERT_VALUES_EQUAL(tables[0].GetDeletes().GetBytes(), 150); + } + + Y_UNIT_TEST(KeepsMissingOptionalFieldsEmpty) { + const TExecStats stats(Ydb::TableStats::QueryStats{}); + UNIT_ASSERT(!stats.GetCompilation()); + UNIT_ASSERT(!stats.GetPlan()); + UNIT_ASSERT(!stats.GetAst()); + UNIT_ASSERT(!stats.GetMeta()); + UNIT_ASSERT(stats.GetQueryPhases().empty()); + } +} diff --git a/tests/unit/client/retry_range/retry_range_ut.cpp b/tests/unit/client/retry_range/retry_range_ut.cpp new file mode 100644 index 00000000000..60dd6b8f1c3 --- /dev/null +++ b/tests/unit/client/retry_range/retry_range_ut.cpp @@ -0,0 +1,373 @@ +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +#include + +#include + +#include +#include + +#include +#include +#include + +#include +#include + +using namespace NYdb; +using namespace NYdb::NStatusHelpers; + +using NYdb::NRetry::TRetryOperationSettings; + +namespace { + +TRetryOperationSettings FastRetrySettings() { + return TRetryOperationSettings() + .MaxRetries(3) + .FastBackoffSettings( + TRetryOperationSettings::DefaultFastBackoffSettings() + .SlotDuration(TDuration::MilliSeconds(1)) + .Ceiling(1)); +} + +TStatus OkStatus() { + return TStatus(EStatus::SUCCESS, NYdb::NIssue::TIssues{}); +} + +TStatus UnavailableStatus() { + return TStatus(EStatus::UNAVAILABLE, NYdb::NIssue::TIssues{}); +} + +TResultSet MakeSingleInt32ResultSet(int32_t value) { + const std::string text = + "columns {\n" + " name: \"v\"\n" + " type { type_id: INT32 }\n" + "}\n" + "rows {\n" + " items { int32_value: " + + std::to_string(value) + + " }\n" + "}\n"; + Ydb::ResultSet proto; + google::protobuf::TextFormat::ParseFromString(TStringType{text}, &proto); + return TResultSet(std::move(proto)); +} + +template +struct TMockStreamIterator { + std::deque> Queue; + + NThreading::TFuture ReadNext() { + UNIT_ASSERT(!Queue.empty()); + NThreading::TFuture f = std::move(Queue.front()); + Queue.pop_front(); + return f; + } + + void Push(Part&& part) { + Queue.push_back(NThreading::MakeFuture(std::move(part))); + } +}; + +void SimulateScanDrainFailure() { + TMockStreamIterator it; + it.Push(NTable::TScanQueryPart(UnavailableStatus())); + NYdb::NRowRangesDetail::DrainStreamIterator(it); +} + +class TMockTableService : public Ydb::Table::V1::TableService::Service { +public: + grpc::Status CreateSession( + grpc::ServerContext*, + const Ydb::Table::CreateSessionRequest*, + Ydb::Table::CreateSessionResponse* response) override + { + Ydb::Table::CreateSessionResult result; + result.set_session_id("fake-table-session-id"); + + auto* op = response->mutable_operation(); + op->set_ready(true); + op->set_status(Ydb::StatusIds::SUCCESS); + op->mutable_result()->PackFrom(result); + return grpc::Status::OK; + } +}; + +class TMockQueryService : public Ydb::Query::V1::QueryService::Service { +public: + grpc::Status CreateSession( + grpc::ServerContext*, + const Ydb::Query::CreateSessionRequest*, + Ydb::Query::CreateSessionResponse* response) override + { + response->set_status(Ydb::StatusIds::SUCCESS); + response->set_session_id("fake-query-session-id"); + response->set_node_id(1); + return grpc::Status::OK; + } + + grpc::Status AttachSession( + grpc::ServerContext*, + const Ydb::Query::AttachSessionRequest*, + grpc::ServerWriter* writer) override + { + Ydb::Query::SessionState state; + state.set_status(Ydb::StatusIds::SUCCESS); + writer->Write(state); + return grpc::Status::OK; + } +}; + +template +std::unique_ptr StartGrpcServer(const std::string& address, TService& service) { + return grpc::ServerBuilder() + .AddListeningPort(TString{address}, grpc::InsecureServerCredentials()) + .RegisterService(&service) + .BuildAndStart(); +} + +struct TTableClientFixture { + TMockTableService TableService; + std::unique_ptr GrpcServer; + std::unique_ptr Driver; + std::unique_ptr Client; + + TTableClientFixture() { + NTesting::InitPortManagerFromEnv(); + const auto portHolder = NTesting::GetFreePort(); + const ui16 port = portHolder; + const auto endpoint = TStringBuilder() << "127.0.0.1:" << port; + + GrpcServer = StartGrpcServer(endpoint, TableService); + Driver = std::make_unique( + TDriverConfig() + .SetEndpoint(endpoint) + .SetDiscoveryMode(EDiscoveryMode::Off) + .SetDatabase("/Root/My/DB")); + Client = std::make_unique(*Driver); + } +}; + +struct TQueryClientFixture { + TMockQueryService QueryService; + std::unique_ptr GrpcServer; + std::unique_ptr Driver; + std::unique_ptr Client; + + TQueryClientFixture() { + NTesting::InitPortManagerFromEnv(); + const auto portHolder = NTesting::GetFreePort(); + const ui16 port = portHolder; + const auto endpoint = TStringBuilder() << "127.0.0.1:" << port; + + GrpcServer = StartGrpcServer(endpoint, QueryService); + Driver = std::make_unique( + TDriverConfig() + .SetEndpoint(endpoint) + .SetDiscoveryMode(EDiscoveryMode::Off) + .SetDatabase("/Root/My/DB")); + Client = std::make_unique(*Driver); + } +}; + +} // namespace + +Y_UNIT_TEST_SUITE(TTableRangeErrorRetryTest) { + Y_UNIT_TEST(RetryOperationSyncWithoutSessionRetriesAfterRangeDrainFailure) { + TTableClientFixture fixture; + size_t attempts = 0; + + auto status = fixture.Client->RetryOperationSync( + [&](NTable::TTableClient&) -> TStatus { + ++attempts; + if (attempts == 1) { + SimulateScanDrainFailure(); + } + return OkStatus(); + }, + FastRetrySettings()); + + UNIT_ASSERT(status.IsSuccess()); + UNIT_ASSERT_VALUES_EQUAL(attempts, 2u); + } + + Y_UNIT_TEST(RetryOperationSyncWithoutSessionPropagatesYdbErrorException) { + TTableClientFixture fixture; + + UNIT_ASSERT_EXCEPTION( + fixture.Client->RetryOperationSync( + [&](NTable::TTableClient&) -> TStatus { + throw TYdbErrorException(UnavailableStatus()); + return OkStatus(); + }, + FastRetrySettings()), + TYdbErrorException); + } + + Y_UNIT_TEST(RetryOperationSyncWithSessionRetriesAfterRangeDrainFailure) { + TTableClientFixture fixture; + size_t attempts = 0; + + auto status = fixture.Client->RetryOperationSync( + [&](NTable::TSession) -> TStatus { + ++attempts; + if (attempts == 1) { + SimulateScanDrainFailure(); + } + return OkStatus(); + }, + FastRetrySettings()); + + UNIT_ASSERT(status.IsSuccess()); + UNIT_ASSERT_VALUES_EQUAL(attempts, 2u); + } + + Y_UNIT_TEST(RetryOperationSyncWithoutSessionRecoversWithResultSetRange) { + TTableClientFixture fixture; + size_t attempts = 0; + int32_t sum = 0; + + auto status = fixture.Client->RetryOperationSync( + [&](NTable::TTableClient&) -> TStatus { + ++attempts; + if (attempts == 1) { + SimulateScanDrainFailure(); + } + + for (auto& row : TRowRange(MakeSingleInt32ResultSet(5))) { + sum += static_cast(row.ColumnParser("v").GetInt32()); + } + return OkStatus(); + }, + FastRetrySettings()); + + UNIT_ASSERT(status.IsSuccess()); + UNIT_ASSERT_VALUES_EQUAL(attempts, 2u); + UNIT_ASSERT_VALUES_EQUAL(sum, 5); + } + + Y_UNIT_TEST(RetryOperationAsyncWithoutSessionRetriesAfterRangeDrainFailure) { + TTableClientFixture fixture; + size_t attempts = 0; + + auto status = fixture.Client->RetryOperation( + [&](NTable::TTableClient&) -> TAsyncStatus { + ++attempts; + if (attempts == 1) { + SimulateScanDrainFailure(); + } + return NThreading::MakeFuture(OkStatus()); + }, + FastRetrySettings()).GetValueSync(); + + UNIT_ASSERT(status.IsSuccess()); + UNIT_ASSERT_VALUES_EQUAL(attempts, 2u); + } + + Y_UNIT_TEST(RetryOperationAsyncWithSessionRetriesAfterRangeDrainFailure) { + TTableClientFixture fixture; + size_t attempts = 0; + + auto status = fixture.Client->RetryOperation( + [&](NTable::TSession) -> TAsyncStatus { + ++attempts; + if (attempts == 1) { + SimulateScanDrainFailure(); + } + return NThreading::MakeFuture(OkStatus()); + }, + FastRetrySettings()).GetValueSync(); + + UNIT_ASSERT(status.IsSuccess()); + UNIT_ASSERT_VALUES_EQUAL(attempts, 2u); + } +} + +Y_UNIT_TEST_SUITE(TQueryRangeErrorRetryTest) { + Y_UNIT_TEST(RetryQuerySyncRetriesAfterRangeDrainFailure) { + TQueryClientFixture fixture; + size_t attempts = 0; + + auto status = fixture.Client->RetryQuerySync( + [&](NQuery::TSession) -> TStatus { + ++attempts; + if (attempts == 1) { + SimulateScanDrainFailure(); + } + return OkStatus(); + }, + FastRetrySettings()); + + UNIT_ASSERT(status.IsSuccess()); + UNIT_ASSERT_VALUES_EQUAL(attempts, 2u); + } + + Y_UNIT_TEST(RetryQuerySyncPropagatesYdbErrorException) { + TQueryClientFixture fixture; + + UNIT_ASSERT_EXCEPTION( + fixture.Client->RetryQuerySync( + [&](NQuery::TSession) -> TStatus { + throw TYdbErrorException(UnavailableStatus()); + return OkStatus(); + }, + FastRetrySettings()), + TYdbErrorException); + } + + Y_UNIT_TEST(RetryQuerySyncPropagatesNonYdbException) { + TQueryClientFixture fixture; + + UNIT_ASSERT_EXCEPTION( + fixture.Client->RetryQuerySync( + [&](NQuery::TSession) -> TStatus { + throw std::runtime_error("user bug"); + return OkStatus(); + }, + FastRetrySettings()), + std::runtime_error); + } + + Y_UNIT_TEST(RetryQueryAsyncRetriesAfterRangeDrainFailure) { + TQueryClientFixture fixture; + size_t attempts = 0; + + auto status = fixture.Client->RetryQuery( + [&](NQuery::TSession) -> TAsyncStatus { + ++attempts; + if (attempts == 1) { + SimulateScanDrainFailure(); + } + return NThreading::MakeFuture(OkStatus()); + }, + FastRetrySettings()).GetValueSync(); + + UNIT_ASSERT(status.IsSuccess()); + UNIT_ASSERT_VALUES_EQUAL(attempts, 2u); + } + + Y_UNIT_TEST(RetryQueryAsyncPropagatesYdbErrorException) { + TQueryClientFixture fixture; + + UNIT_ASSERT_EXCEPTION( + fixture.Client->RetryQuery( + [&](NQuery::TSession) -> TAsyncStatus { + throw TYdbErrorException(UnavailableStatus()); + return NThreading::MakeFuture(OkStatus()); + }, + FastRetrySettings()).GetValueSync(), + TYdbErrorException); + } +} diff --git a/tests/unit/client/row_ranges/row_ranges_fixtures.h b/tests/unit/client/row_ranges/row_ranges_fixtures.h new file mode 100644 index 00000000000..766a7803ecf --- /dev/null +++ b/tests/unit/client/row_ranges/row_ranges_fixtures.h @@ -0,0 +1,3 @@ +#pragma once + +#include "row_ranges_test_helpers.h" diff --git a/tests/unit/client/row_ranges/row_ranges_public_include_ut.cpp b/tests/unit/client/row_ranges/row_ranges_public_include_ut.cpp new file mode 100644 index 00000000000..02620e934f8 --- /dev/null +++ b/tests/unit/client/row_ranges/row_ranges_public_include_ut.cpp @@ -0,0 +1,54 @@ +#include + +#include + +#include "row_ranges_test_helpers.h" + +#include +#include +#include +#include +#include +#include + +using namespace NYdb; +namespace { + +static_assert(std::is_same_v< + std::iterator_traits::iterator_category, + std::input_iterator_tag>); +static_assert(!std::copy_constructible); +static_assert(!std::assignable_from); +static_assert(std::movable); +static_assert(std::sentinel_for); + +using Int32Columns = decltype( + std::declval().Get(std::declval>())); +using TupleColumns = decltype( + std::declval().Get>( + std::declval>())); + +static_assert(std::same_as< + typename Int32Columns::Iterator::value_type, + std::tuple>); +static_assert(std::same_as< + typename TupleColumns::Iterator::value_type, + std::tuple>>); + +} // namespace + +Y_UNIT_TEST_SUITE(TRowRangePublicIncludeTest) { + Y_UNIT_TEST(SmokeIterateAndGet) { + TRowRange range(NRowRangesTest::MakeSingleInt32ResultSet(42)); + int32_t sum = 0; + for (auto it = range.begin(); it != range.end(); ++it) { + sum += static_cast(it->ColumnParser("v").GetInt32()); + } + UNIT_ASSERT_VALUES_EQUAL(sum, 42); + + TRowRange range2(NRowRangesTest::MakeSingleInt32ResultSet(42)); + for (auto [v] : range2.Get({"v"})) { + UNIT_ASSERT_VALUES_EQUAL(v, 42); + } + } +} diff --git a/tests/unit/client/row_ranges/row_ranges_test_helpers.h b/tests/unit/client/row_ranges/row_ranges_test_helpers.h new file mode 100644 index 00000000000..a8528276859 --- /dev/null +++ b/tests/unit/client/row_ranges/row_ranges_test_helpers.h @@ -0,0 +1,333 @@ +#pragma once + +#include +#include +#include +#include + +#include + +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace NRowRangesTest { + +inline NYdb::TType PrimitiveType(NYdb::EPrimitiveType primitive) { + return NYdb::TTypeBuilder().Primitive(primitive).Build(); +} + +class TResultSetBuilder { +public: + TResultSetBuilder& Column(std::string name, const NYdb::TType& type) { + auto* column = Proto_.add_columns(); + column->set_name(std::move(name)); + column->mutable_type()->CopyFrom(type.GetProto()); + return *this; + } + + TResultSetBuilder& Row(const NYdb::TValue& cell) { + return Row(std::vector{cell}); + } + + TResultSetBuilder& Row(std::vector cells) { + auto* row = Proto_.add_rows(); + for (const auto& cell : cells) { + row->add_items()->CopyFrom(cell.GetProto()); + } + return *this; + } + + NYdb::TResultSet Build() { + return NYdb::TResultSet(std::move(Proto_)); + } + +private: + Ydb::ResultSet Proto_; +}; + +inline NYdb::TResultSet SingleColumn(std::string name, NYdb::TValue value) { + return TResultSetBuilder() + .Column(std::move(name), value.GetType()) + .Row(std::move(value)) + .Build(); +} + +inline NYdb::TResultSet MakeSingleInt32ResultSet(int32_t value) { + return SingleColumn("v", NYdb::TValueBuilder().Int32(value).Build()); +} + +inline NYdb::TResultSet MakeEmptyInt32SchemaResultSet() { + return TResultSetBuilder() + .Column("v", PrimitiveType(NYdb::EPrimitiveType::Int32)) + .Build(); +} + +inline NYdb::TResultSet MakeThreeInt32RowsResultSet() { + TResultSetBuilder builder; + builder.Column("v", PrimitiveType(NYdb::EPrimitiveType::Int32)); + for (int32_t value : {1, 2, 3}) { + builder.Row(NYdb::TValueBuilder().Int32(value).Build()); + } + return builder.Build(); +} + +inline NYdb::TResultSet MakeOptionalInt32ResultSet(const std::vector>& values) { + TResultSetBuilder builder; + builder.Column("v", NYdb::TTypeBuilder() + .BeginOptional() + .Primitive(NYdb::EPrimitiveType::Int32) + .EndOptional() + .Build()); + for (const auto& value : values) { + if (value) { + builder.Row(NYdb::TValueBuilder().OptionalInt32(value).Build()); + } else { + builder.Row(NYdb::TValueBuilder().EmptyOptional(NYdb::EPrimitiveType::Int32).Build()); + } + } + return builder.Build(); +} + +inline NYdb::TResultSet MakeIdNameResultSet(const std::vector>& rows) { + TResultSetBuilder builder; + builder.Column("id", PrimitiveType(NYdb::EPrimitiveType::Int32)); + builder.Column("name", PrimitiveType(NYdb::EPrimitiveType::Utf8)); + for (const auto& [id, name] : rows) { + builder.Row({ + NYdb::TValueBuilder().Int32(id).Build(), + NYdb::TValueBuilder().Utf8(name).Build(), + }); + } + return builder.Build(); +} + +inline NYdb::TResultSet MakeSingleDateResultSet(uint32_t daysSinceEpoch) { + const TInstant instant = TInstant::Seconds(static_cast(daysSinceEpoch) * 86400ULL); + return SingleColumn("d", NYdb::TValueBuilder().Date(instant).Build()); +} + +inline NYdb::TResultSet MakeSingleDatetimeResultSet(uint32_t secondsSinceEpoch) { + return SingleColumn( + "d", + NYdb::TValueBuilder().Datetime(TInstant::Seconds(secondsSinceEpoch)).Build()); +} + +inline NYdb::TResultSet MakeSingleTimestampResultSet(uint64_t microsSinceEpoch) { + return SingleColumn( + "d", + NYdb::TValueBuilder().Timestamp(TInstant::MicroSeconds(microsSinceEpoch)).Build()); +} + +inline NYdb::TResultSet MakeOptionalDateResultSet(const std::vector>& values) { + TResultSetBuilder builder; + builder.Column("d", NYdb::TTypeBuilder() + .BeginOptional() + .Primitive(NYdb::EPrimitiveType::Date) + .EndOptional() + .Build()); + for (const auto& value : values) { + if (value) { + const TInstant instant = TInstant::Seconds(static_cast(*value) * 86400ULL); + builder.Row(NYdb::TValueBuilder().OptionalDate(instant).Build()); + } else { + builder.Row(NYdb::TValueBuilder().EmptyOptional(NYdb::EPrimitiveType::Date).Build()); + } + } + return builder.Build(); +} + +inline NYdb::TResultSet MakeSingleIntervalResultSet(int64_t micros) { + return SingleColumn("v", NYdb::TValueBuilder().Interval(micros).Build()); +} + +inline NYdb::TResultSet MakeSingleJsonResultSet(const std::string& json) { + return SingleColumn("v", NYdb::TValueBuilder().Json(json).Build()); +} + +inline NYdb::TResultSet MakeSingleBytesResultSet(const std::string& bytesLiteral) { + return SingleColumn("v", NYdb::TValueBuilder().String(bytesLiteral).Build()); +} + +inline NYdb::TResultSet MakeSingleTimestamp64ResultSet(int64_t micros) { + const auto value = std::chrono::sys_time(NYdb::TWideMicroseconds(micros)); + return SingleColumn("v", NYdb::TValueBuilder().Timestamp64(value).Build()); +} + +inline NYdb::TResultSet MakeSingleInterval64ResultSet(int64_t micros) { + return SingleColumn("v", NYdb::TValueBuilder().Interval64(NYdb::TWideMicroseconds(micros)).Build()); +} + +inline NYdb::TResultSet MakeListInt32ResultSet(const std::vector& values) { + NYdb::TValueBuilder list; + list.BeginList(); + for (int32_t value : values) { + list.AddListItem().Int32(value); + } + list.EndList(); + return SingleColumn("v", list.Build()); +} + +inline NYdb::TResultSet MakeTupleColumnResultSet() { + const NYdb::TValue value = NYdb::TValueBuilder() + .BeginTuple() + .AddElement().BeginOptional().Utf8("hello").EndOptional() + .AddElement().Int8(-5) + .EndTuple() + .Build(); + return SingleColumn("t", value); +} + +inline NYdb::TResultSet MakeStructColumnResultSet() { + const NYdb::TValue value = NYdb::TValueBuilder() + .BeginStruct() + .AddMember("id").Int32(42) + .AddMember("name").Utf8("struct") + .EndStruct() + .Build(); + return SingleColumn("s", value); +} + +inline NYdb::TResultSet MakeDictColumnResultSet() { + const NYdb::TValue value = NYdb::TValueBuilder() + .BeginDict() + .AddDictItem().DictKey().Uint32(1).DictPayload().Utf8("one") + .AddDictItem().DictKey().Uint32(2).DictPayload().Utf8("two") + .EndDict() + .Build(); + return SingleColumn("d", value); +} + +inline NYdb::TResultSet MakeOptionalListInt32ResultSet() { + const std::string text = + "columns {\n" + " name: \"v\"\n" + " type { optional_type { item { list_type { item { type_id: INT32 } } } } }\n" + "}\n" + "rows {\n" + " items { null_flag_value: NULL_VALUE }\n" + "}\n" + "rows {\n" + " items {\n" + " items { int32_value: 5 }\n" + " items { int32_value: 6 }\n" + " }\n" + "}\n"; + Ydb::ResultSet proto; + google::protobuf::TextFormat::ParseFromString(text, &proto); + return NYdb::TResultSet(std::move(proto)); +} + +inline NYdb::TResultSet MakeTaggedColumnResultSet() { + const NYdb::TValue value = NYdb::TValueBuilder() + .BeginTagged("my_tag") + .Utf8("tagged_value") + .EndTagged() + .Build(); + return SingleColumn("t", value); +} + +inline NYdb::TValue MakeVariantUtf8Value() { + Ydb::Type type; + auto* tupleItems = type.mutable_variant_type()->mutable_tuple_items(); + tupleItems->add_elements()->set_type_id(Ydb::Type::INT32); + tupleItems->add_elements()->set_type_id(Ydb::Type::UTF8); + Ydb::Value value; + value.set_variant_index(1); + value.mutable_nested_value()->set_text_value("variant_utf8"); + return NYdb::TValue(NYdb::TType(type), std::move(value)); +} + +inline NYdb::TResultSet MakeVariantColumnResultSet() { + return SingleColumn("v", MakeVariantUtf8Value()); +} + +inline NYdb::TResultSet MakeSingleColumnResultSet(const std::string& columnName, const NYdb::TValue& value) { + return SingleColumn(columnName, value); +} + +inline int32_t SumColumnV(const NYdb::TResultSet& rs) { + int32_t sum = 0; + NYdb::TResultSetParser parser(rs); + while (parser.TryNextRow()) { + sum += static_cast(parser.ColumnParser("v").GetInt32()); + } + return sum; +} + +inline NYdb::NTable::TDataQueryResult MakeDataQueryResult(std::vector&& sets) { + return NYdb::NTable::TDataQueryResult( + NYdb::TStatus(NYdb::EStatus::SUCCESS, NYdb::NIssue::TIssues{}), + std::move(sets), + std::nullopt, + std::nullopt, + false, + std::nullopt); +} + +template +T ReadSingle(NYdb::TResultSet&& rs, std::string_view col) { + NYdb::TRowRange range(std::move(rs)); + for (auto [value] : range.Get({col})) { + return value; + } + ythrow yexception() << "empty result set"; +} + +template +std::vector ReadAll(NYdb::TResultSet&& rs, std::string_view col) { + std::vector result; + NYdb::TRowRange range(std::move(rs)); + for (auto [value] : range.Get({col})) { + result.push_back(std::move(value)); + } + return result; +} + +template +std::tuple ReadSingleRow( + NYdb::TResultSet&& rs, + std::initializer_list cols) { + NYdb::TRowRange range(std::move(rs)); + for (auto row : range.Get(cols)) { + return row; + } + ythrow yexception() << "empty result set"; +} + +template +std::vector> ReadAllRows( + NYdb::TResultSet&& rs, + std::initializer_list cols) { + std::vector> result; + NYdb::TRowRange range(std::move(rs)); + for (auto row : range.Get(cols)) { + result.push_back(std::move(row)); + } + return result; +} + +template +void AssertGetThrows(NYdb::TResultSet&& rs, std::initializer_list cols) { + NYdb::TRowRange range(std::move(rs)); + UNIT_ASSERT_EXCEPTION( + ([&] { + for (auto row : range.Get(cols)) { + (void)row; + } + }()), + Exception); +} + +} // namespace NRowRangesTest diff --git a/tests/unit/client/row_ranges/row_ranges_ut.cpp b/tests/unit/client/row_ranges/row_ranges_ut.cpp new file mode 100644 index 00000000000..be4b6dfc890 --- /dev/null +++ b/tests/unit/client/row_ranges/row_ranges_ut.cpp @@ -0,0 +1,386 @@ +#include +#include +#include +#include + +#include "row_ranges_test_helpers.h" + +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace NYdb; +using namespace NYdb::NStatusHelpers; +using namespace NRowRangesTest; + +namespace { + +TStatus OkStatus() { + return TStatus(EStatus::SUCCESS, NYdb::NIssue::TIssues{}); +} + +TStatus EosStatus() { + return TStatus(EStatus::CLIENT_OUT_OF_RANGE, NYdb::NIssue::TIssues{}); +} + +template +struct TMockStreamIterator { + std::deque> Queue; + + NThreading::TFuture ReadNext() { + UNIT_ASSERT(!Queue.empty()); + NThreading::TFuture f = std::move(Queue.front()); + Queue.pop_front(); + return f; + } + + void Push(Part&& part) { + Queue.push_back(NThreading::MakeFuture(std::move(part))); + } +}; + +static_assert(std::is_same_v< + std::iterator_traits::iterator_category, + std::input_iterator_tag>); +static_assert(!std::copy_constructible); +static_assert(!std::assignable_from); +static_assert(std::movable); +static_assert(std::sentinel_for); + +using ListInt32Columns = decltype(std::declval().Get>({"v"})); +using DequeInt32Columns = decltype(std::declval().Get>({"v"})); +static_assert(std::same_as< + typename ListInt32Columns::Iterator::value_type, + std::tuple>>); +static_assert(std::same_as< + typename DequeInt32Columns::Iterator::value_type, + std::tuple>>); + +} // namespace + +Y_UNIT_TEST_SUITE(TRowRangeTest) { + Y_UNIT_TEST(EmptyDataQueryResultYieldsNoRows) { + TRowRange range(MakeDataQueryResult({})); + size_t count = 0; + for (auto it = range.begin(); it != range.end(); ++it) { + (void)it; + ++count; + } + UNIT_ASSERT_VALUES_EQUAL(count, 0u); + } + + Y_UNIT_TEST(SingleResultSetCtorIteratesRows) { + TRowRange range(MakeSingleInt32ResultSet(7)); + int sum = 0; + for (auto it = range.begin(); it != range.end(); ++it) { + sum += static_cast(it->ColumnParser("v").GetInt32()); + } + UNIT_ASSERT_VALUES_EQUAL(sum, 7); + } + + Y_UNIT_TEST(MultiSetDataQueryResultThrows) { + std::vector sets; + sets.push_back(MakeSingleInt32ResultSet(1)); + sets.push_back(MakeSingleInt32ResultSet(2)); + sets.push_back(MakeSingleInt32ResultSet(3)); + UNIT_ASSERT_EXCEPTION(TRowRange(MakeDataQueryResult(std::move(sets))), yexception); + } + + Y_UNIT_TEST(FailedDataQueryResultCtorThrows) { + auto data = NTable::TDataQueryResult( + TStatus(EStatus::UNAVAILABLE, NYdb::NIssue::TIssues{}), + {}, + std::nullopt, + std::nullopt, + false, + std::nullopt); + UNIT_ASSERT_EXCEPTION(TRowRange(std::move(data)), TYdbRangeErrorException); + } +} + +Y_UNIT_TEST_SUITE(TRowRangeStreamDrainTest) { + Y_UNIT_TEST(DrainSkipsNonResultParts) { + { + TMockStreamIterator it; + it.Push(NQuery::TExecuteQueryPart(OkStatus(), std::nullopt, std::nullopt)); + it.Push(NQuery::TExecuteQueryPart(OkStatus(), MakeSingleInt32ResultSet(11), 0, std::nullopt, std::nullopt)); + it.Push(NQuery::TExecuteQueryPart(EosStatus(), std::nullopt, std::nullopt)); + UNIT_ASSERT_VALUES_EQUAL(SumColumnV(*NYdb::NRowRangesDetail::DrainStreamIterator(it)), 11); + UNIT_ASSERT(!NYdb::NRowRangesDetail::DrainStreamIterator(it).has_value()); + } + { + TMockStreamIterator it; + it.Push(NTable::TScanQueryPart(OkStatus())); + it.Push(NTable::TScanQueryPart(OkStatus(), MakeSingleInt32ResultSet(22), std::nullopt, std::nullopt, std::nullopt)); + it.Push(NTable::TScanQueryPart(EosStatus())); + UNIT_ASSERT_VALUES_EQUAL(SumColumnV(*NYdb::NRowRangesDetail::DrainStreamIterator(it)), 22); + UNIT_ASSERT(!NYdb::NRowRangesDetail::DrainStreamIterator(it).has_value()); + } + { + TMockStreamIterator it; + it.Push(NTable::TReadTableResultPart(MakeSingleInt32ResultSet(33), OkStatus())); + it.Push(NTable::TReadTableResultPart(MakeEmptyInt32SchemaResultSet(), EosStatus())); + UNIT_ASSERT_VALUES_EQUAL(SumColumnV(*NYdb::NRowRangesDetail::DrainStreamIterator(it)), 33); + UNIT_ASSERT(!NYdb::NRowRangesDetail::DrainStreamIterator(it).has_value()); + } + } + + Y_UNIT_TEST(ExecuteQueryStreamSameIndexPartsOk) { + TMockStreamIterator it; + it.Push(NQuery::TExecuteQueryPart(OkStatus(), MakeSingleInt32ResultSet(1), 0, std::nullopt, std::nullopt)); + it.Push(NQuery::TExecuteQueryPart(OkStatus(), MakeSingleInt32ResultSet(2), 0, std::nullopt, std::nullopt)); + it.Push(NQuery::TExecuteQueryPart(EosStatus(), std::nullopt, std::nullopt)); + + int32_t sum = 0; + while (auto rs = NYdb::NRowRangesDetail::DrainStreamIterator(it)) { + sum += SumColumnV(*rs); + } + UNIT_ASSERT_VALUES_EQUAL(sum, 3); + } + + Y_UNIT_TEST(ExecuteQueryStreamNonZeroIndexThrows) { + TMockStreamIterator it; + it.Push(NQuery::TExecuteQueryPart(OkStatus(), MakeSingleInt32ResultSet(1), 1, std::nullopt, std::nullopt)); + UNIT_ASSERT_EXCEPTION(NYdb::NRowRangesDetail::DrainStreamIterator(it), yexception); + } +} + +Y_UNIT_TEST_SUITE(TRowColumnsTest) { + Y_UNIT_TEST(GetSingleInt32) { + UNIT_ASSERT_VALUES_EQUAL(ReadSingle(MakeSingleInt32ResultSet(7), "v"), 7); + } + + Y_UNIT_TEST(GetTwoColumnsTuple) { + auto got = ReadAllRows( + MakeIdNameResultSet({{1, "foo"}, {2, "bar"}, {3, "baz"}}), + {"id", "name"}); + UNIT_ASSERT_VALUES_EQUAL(got.size(), 3u); + UNIT_ASSERT_VALUES_EQUAL(std::get<0>(got[0]), 1); + UNIT_ASSERT_VALUES_EQUAL(std::get<1>(got[0]), "foo"); + UNIT_ASSERT_VALUES_EQUAL(std::get<0>(got[2]), 3); + UNIT_ASSERT_VALUES_EQUAL(std::get<1>(got[2]), "baz"); + } + + Y_UNIT_TEST(GetOptional) { + auto got = ReadAll>( + MakeOptionalInt32ResultSet({std::nullopt, 42, std::nullopt, 7}), + "v"); + UNIT_ASSERT_VALUES_EQUAL(got.size(), 4u); + UNIT_ASSERT(!got[0].has_value()); + UNIT_ASSERT_VALUES_EQUAL(*got[1], 42); + UNIT_ASSERT(!got[2].has_value()); + UNIT_ASSERT_VALUES_EQUAL(*got[3], 7); + } + + Y_UNIT_TEST(GetAcceptsVectorOfNames) { + TRowRange range(MakeIdNameResultSet({{10, "x"}, {20, "y"}})); + std::vector names{"id", "name"}; + int32_t idSum = 0; + std::string concat; + for (auto [id, n] : range.Get(names)) { + idSum += id; + concat += n; + } + UNIT_ASSERT_VALUES_EQUAL(idSum, 30); + UNIT_ASSERT_VALUES_EQUAL(concat, "xy"); + } + + Y_UNIT_TEST(GetColumnNameCountMismatchThrows) { + TRowRange range(MakeSingleInt32ResultSet(1)); + UNIT_ASSERT_EXCEPTION((range.Get({"v"})), std::invalid_argument); + UNIT_ASSERT_EXCEPTION((range.Get({"v", "extra"})), std::invalid_argument); + } + + Y_UNIT_TEST(GetMultipleResultSetsThrows) { + std::vector sets; + sets.push_back(MakeIdNameResultSet({{1, "a"}})); + sets.push_back(MakeIdNameResultSet({{2, "b"}, {3, "c"}})); + sets.push_back(MakeIdNameResultSet({{4, "d"}})); + UNIT_ASSERT_EXCEPTION(TRowRange(MakeDataQueryResult(std::move(sets))), yexception); + } + + Y_UNIT_TEST(GetTInstantDispatchesPrimitive) { + constexpr uint32_t days = 18993; + UNIT_ASSERT_VALUES_EQUAL(ReadSingle(MakeSingleDateResultSet(days), "d").Days(), days); + + constexpr uint32_t secs = 1640995200u; + UNIT_ASSERT_VALUES_EQUAL(ReadSingle(MakeSingleDatetimeResultSet(secs), "d").Seconds(), secs); + + constexpr uint64_t micros = 1640995200000123ull; + UNIT_ASSERT_VALUES_EQUAL(ReadSingle(MakeSingleTimestampResultSet(micros), "d").MicroSeconds(), micros); + } + + Y_UNIT_TEST(GetOptionalTInstant) { + constexpr uint32_t days = 18993; + auto got = ReadAll>( + MakeOptionalDateResultSet({std::nullopt, days, std::nullopt}), + "d"); + UNIT_ASSERT_VALUES_EQUAL(got.size(), 3u); + UNIT_ASSERT(!got[0].has_value()); + UNIT_ASSERT(got[1].has_value()); + UNIT_ASSERT_VALUES_EQUAL(got[1]->Days(), days); + UNIT_ASSERT(!got[2].has_value()); + } + + Y_UNIT_TEST(GetTInstantWrongTypeThrows) { + AssertGetThrows(MakeSingleInt32ResultSet(1), {"v"}); + } + + Y_UNIT_TEST(GetStringLikePrimitives) { + UNIT_ASSERT_VALUES_EQUAL(ReadSingle(MakeSingleJsonResultSet("[1,2,3]"), "v"), "[1,2,3]"); + UNIT_ASSERT_VALUES_EQUAL(ReadSingle(MakeSingleBytesResultSet("raw_bytes"), "v"), "raw_bytes"); + } + + Y_UNIT_TEST(GetWideTimeScalars) { + constexpr int64_t intervalMicros = 123456789; + UNIT_ASSERT_VALUES_EQUAL(ReadSingle(MakeSingleIntervalResultSet(intervalMicros), "v"), intervalMicros); + + constexpr int64_t timestamp64Micros = 1640995200000123LL; + const auto timestamp64 = ReadSingle>( + MakeSingleTimestamp64ResultSet(timestamp64Micros), + "v"); + UNIT_ASSERT_VALUES_EQUAL(timestamp64.time_since_epoch().count(), timestamp64Micros); + + constexpr int64_t interval64Micros = 9876543210LL; + UNIT_ASSERT_VALUES_EQUAL( + ReadSingle(MakeSingleInterval64ResultSet(interval64Micros), "v").count(), + interval64Micros); + } + + Y_UNIT_TEST(GetDecimalAndPg) { + auto decimalCell = TValueBuilder().Decimal(TDecimalValue("42.50", 10, 2)).Build(); + UNIT_ASSERT_VALUES_EQUAL( + ReadSingle(MakeSingleColumnResultSet("d", decimalCell), "d").ToString(), + "42.5"); + + auto pgCell = TValueBuilder() + .Pg(TPgValue(TPgValue::VK_TEXT, "pg_text", TPgType("text"))) + .Build(); + const auto pg = ReadSingle(MakeSingleColumnResultSet("p", pgCell), "p"); + UNIT_ASSERT(pg.IsText()); + UNIT_ASSERT_VALUES_EQUAL(pg.Content_, "pg_text"); + } + + Y_UNIT_TEST(GetTValueColumn) { + auto cell = TValueBuilder().Int32(99).Build(); + const auto value = ReadSingle(MakeSingleColumnResultSet("v", cell), "v"); + UNIT_ASSERT_VALUES_EQUAL(TValueParser(value).GetInt32(), 99); + } + + Y_UNIT_TEST(GetSequenceContainers) { + const auto got = ReadSingle>(MakeListInt32ResultSet({1, 2, 3}), "v"); + UNIT_ASSERT_VALUES_EQUAL(got.size(), 3u); + UNIT_ASSERT_VALUES_EQUAL(got[0], 1); + UNIT_ASSERT_VALUES_EQUAL(got[2], 3); + } + + Y_UNIT_TEST(GetAssociativeContainers) { + const auto check = [](auto map) { + UNIT_ASSERT_VALUES_EQUAL(map.size(), 2u); + UNIT_ASSERT_VALUES_EQUAL(map.at(1), "one"); + UNIT_ASSERT_VALUES_EQUAL(map.at(2), "two"); + }; + check(ReadSingle>(MakeDictColumnResultSet(), "d")); + check(ReadSingle>(MakeDictColumnResultSet(), "d")); + } + + Y_UNIT_TEST(GetTupleColumn) { + const auto got = ReadSingle, int8_t>>(MakeTupleColumnResultSet(), "t"); + UNIT_ASSERT(std::get<0>(got).has_value()); + UNIT_ASSERT_VALUES_EQUAL(*std::get<0>(got), "hello"); + UNIT_ASSERT_VALUES_EQUAL(std::get<1>(got), static_cast(-5)); + } + + Y_UNIT_TEST(GetStructAsTuple) { + const auto got = ReadSingle>(MakeStructColumnResultSet(), "s"); + UNIT_ASSERT_VALUES_EQUAL(std::get<0>(got), 42); + UNIT_ASSERT_VALUES_EQUAL(std::get<1>(got), "struct"); + } + + Y_UNIT_TEST(GetTupleArityMismatchThrows) { + AssertGetThrows< + std::tuple, int8_t, int32_t>, + std::runtime_error>(MakeTupleColumnResultSet(), {"t"}); + AssertGetThrows>, std::runtime_error>(MakeTupleColumnResultSet(), {"t"}); + } + + Y_UNIT_TEST(GetStructArityMismatchThrows) { + AssertGetThrows, std::runtime_error>(MakeStructColumnResultSet(), {"s"}); + AssertGetThrows, std::runtime_error>(MakeStructColumnResultSet(), {"s"}); + } + + Y_UNIT_TEST(GetOptionalVectorInt32) { + const auto got = ReadAll>>(MakeOptionalListInt32ResultSet(), "v"); + UNIT_ASSERT_VALUES_EQUAL(got.size(), 2u); + UNIT_ASSERT(!got[0].has_value()); + UNIT_ASSERT(got[1].has_value()); + UNIT_ASSERT_VALUES_EQUAL(got[1]->size(), 2u); + UNIT_ASSERT_VALUES_EQUAL((*got[1])[0], 5); + UNIT_ASSERT_VALUES_EQUAL((*got[1])[1], 6); + } + + Y_UNIT_TEST(GetTaggedColumn) { + const auto got = ReadSingle>(MakeTaggedColumnResultSet(), "t"); + UNIT_ASSERT_VALUES_EQUAL(got.first, "my_tag"); + UNIT_ASSERT_VALUES_EQUAL(got.second, "tagged_value"); + } + + Y_UNIT_TEST(GetVariantColumn) { + const auto got = ReadSingle>(MakeVariantColumnResultSet(), "v"); + UNIT_ASSERT(std::holds_alternative(got)); + UNIT_ASSERT_VALUES_EQUAL(std::get(got), "variant_utf8"); + } +} + +Y_UNIT_TEST_SUITE(TRowIteratorTest) { + Y_UNIT_TEST(StdDistanceMatchesRowCount) { + TRowRange range(MakeThreeInt32RowsResultSet()); + UNIT_ASSERT_VALUES_EQUAL(std::ranges::distance(range.begin(), range.end()), 3); + } + + Y_UNIT_TEST(PrefixAndPostfixIncrement) { + TRowRange range(MakeThreeInt32RowsResultSet()); + auto it = range.begin(); + UNIT_ASSERT_VALUES_EQUAL(static_cast(it->ColumnParser("v").GetInt32()), 1); + ++it; + UNIT_ASSERT_VALUES_EQUAL(static_cast(it->ColumnParser("v").GetInt32()), 2); + it++; + UNIT_ASSERT_VALUES_EQUAL(static_cast(it->ColumnParser("v").GetInt32()), 3); + ++it; + UNIT_ASSERT(it == range.end()); + } + + Y_UNIT_TEST(SurvivesRangeMove) { + TRowRange range(MakeThreeInt32RowsResultSet()); + auto it = range.begin(); + TRowRange moved = std::move(range); + (void)moved; + int32_t sum = 0; + for (; it != TRowIterEnd{}; ++it) { + sum += static_cast(it->ColumnParser("v").GetInt32()); + } + UNIT_ASSERT_VALUES_EQUAL(sum, 6); + } + + Y_UNIT_TEST(TypedViewSurvivesRangeMove) { + TRowRange range(MakeThreeInt32RowsResultSet()); + auto cols = range.Get({"v"}); + TRowRange moved = std::move(range); + (void)moved; + int32_t sum = 0; + for (auto [v] : cols) { + sum += v; + } + UNIT_ASSERT_VALUES_EQUAL(sum, 6); + } +}