|
| 1 | +#include <ydb/core/persqueue/ut/common/autoscaling_ut_common.h> |
| 2 | + |
| 3 | +#include <ydb/public/sdk/cpp/client/ydb_topic/ut/ut_utils/topic_sdk_test_setup.h> |
| 4 | + |
| 5 | +#include <library/cpp/testing/unittest/registar.h> |
| 6 | +#include <ydb/core/persqueue/partition_key_range/partition_key_range.h> |
| 7 | +#include <ydb/core/persqueue/partition_scale_manager.h> |
| 8 | +#include <ydb/core/tx/schemeshard/ut_helpers/helpers.h> |
| 9 | +#include <ydb/core/tx/schemeshard/ut_helpers/test_env.h> |
| 10 | + |
| 11 | +#include <util/stream/output.h> |
| 12 | + |
| 13 | +namespace NKikimr { |
| 14 | + |
| 15 | +using namespace NYdb::NTopic; |
| 16 | +using namespace NYdb::NTopic::NTests; |
| 17 | +using namespace NSchemeShardUT_Private; |
| 18 | +using namespace NKikimr::NPQ::NTest; |
| 19 | + |
| 20 | +Y_UNIT_TEST_SUITE(SlowTopicAutopartitioning) { |
| 21 | + |
| 22 | + void ExecuteQuery(NYdb::NTable::TSession& session, const TString& query ) { |
| 23 | + const auto result = session.ExecuteSchemeQuery(query).GetValueSync(); |
| 24 | + UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString()); |
| 25 | + } |
| 26 | + |
| 27 | + void ExecuteDataQuery(NYdb::NTable::TSession& session, const TString& query ) { |
| 28 | + const auto result = session.ExecuteDataQuery(query, NYdb::NTable::TTxControl::BeginTx().CommitTx()).GetValueSync(); |
| 29 | + UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), NYdb::EStatus::SUCCESS, result.GetIssues().ToString()); |
| 30 | + } |
| 31 | + |
| 32 | + ui64 GetBalancerTabletId(TTopicSdkTestSetup& setup, const TString& topicPath) { |
| 33 | + auto pathDescr = setup.GetServer().AnnoyingClient->Ls(topicPath)->Record.GetPathDescription().GetSelf(); |
| 34 | + auto balancerTabletId = pathDescr.GetBalancerTabletID(); |
| 35 | + Cerr << ">>>>> BalancerTabletID=" << balancerTabletId << Endl << Flush; |
| 36 | + UNIT_ASSERT(balancerTabletId); |
| 37 | + return balancerTabletId; |
| 38 | + } |
| 39 | + |
| 40 | + void SplitPartitionRB(TTopicSdkTestSetup& setup, const TString& topicPath, ui32 partitionId) { |
| 41 | + auto balancerTabletId = GetBalancerTabletId(setup, topicPath); |
| 42 | + auto edge = setup.GetRuntime().AllocateEdgeActor(); |
| 43 | + setup.GetRuntime().SendToPipe(balancerTabletId, edge, new TEvPQ::TEvPartitionScaleStatusChanged(partitionId, NKikimrPQ::EScaleStatus::NEED_SPLIT)); |
| 44 | + } |
| 45 | + |
| 46 | + void AssertPartitionCount(TTopicSdkTestSetup& setup, const TString& topicPath, size_t expectedCount) { |
| 47 | + auto client = setup.MakeClient(); |
| 48 | + auto describe = client.DescribeTopic(topicPath).GetValueSync(); |
| 49 | + UNIT_ASSERT_VALUES_EQUAL(describe.GetTopicDescription().GetPartitions().size(), expectedCount); |
| 50 | + } |
| 51 | + |
| 52 | + void WaitAndAssertPartitionCount(TTopicSdkTestSetup& setup, const TString& topicPath, size_t expectedCount) { |
| 53 | + auto client = setup.MakeClient(); |
| 54 | + size_t partitionCount = 0; |
| 55 | + for (size_t i = 0; i < 10; ++i) { |
| 56 | + Sleep(TDuration::Seconds(1)); |
| 57 | + auto describe = client.DescribeTopic(topicPath).GetValueSync(); |
| 58 | + partitionCount = describe.GetTopicDescription().GetPartitions().size(); |
| 59 | + if (partitionCount == expectedCount) { |
| 60 | + break; |
| 61 | + } |
| 62 | + } |
| 63 | + UNIT_ASSERT_VALUES_EQUAL(partitionCount, expectedCount); |
| 64 | + } |
| 65 | + |
| 66 | + void AssertMessageCountInTopic(TTopicClient client, const TString& topicPath, size_t expected, TDuration timeout = TDuration::Seconds(300)) { |
| 67 | + TInstant deadlineTime = TInstant::Now() + timeout; |
| 68 | + |
| 69 | + size_t count = 0; |
| 70 | + |
| 71 | + auto reader = client.CreateReadSession( |
| 72 | + TReadSessionSettings() |
| 73 | + .AutoPartitioningSupport(true) |
| 74 | + .AppendTopics(TTopicReadSettings(topicPath)) |
| 75 | + .ConsumerName("consumer-1")); |
| 76 | + while(deadlineTime > TInstant::Now()) { |
| 77 | + for (auto event : reader->GetEvents(false)) { |
| 78 | + if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent>(&event)) { |
| 79 | + count += x->GetMessages().size(); |
| 80 | + } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TStartPartitionSessionEvent>(&event)) { |
| 81 | + x->Confirm(); |
| 82 | + Cerr << ">>>>> " << x->DebugString() << Endl << Flush; |
| 83 | + } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TCommitOffsetAcknowledgementEvent>(&event)) { |
| 84 | + Cerr << ">>>>> " << x->DebugString() << Endl << Flush; |
| 85 | + } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TPartitionSessionStatusEvent>(&event)) { |
| 86 | + Cerr << ">>>>> " << x->DebugString() << Endl << Flush; |
| 87 | + } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TStopPartitionSessionEvent>(&event)) { |
| 88 | + x->Confirm(); |
| 89 | + Cerr << ">>>>> " << x->DebugString() << Endl << Flush; |
| 90 | + } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TPartitionSessionClosedEvent>(&event)) { |
| 91 | + Cerr << ">>>>> " << x->DebugString() << Endl << Flush; |
| 92 | + } else if (auto* x = std::get_if<NYdb::NTopic::TReadSessionEvent::TEndPartitionSessionEvent>(&event)) { |
| 93 | + x->Confirm(); |
| 94 | + Cerr << ">>>>> " << x->DebugString() << Endl << Flush; |
| 95 | + } else if (auto* sessionClosedEvent = std::get_if<NYdb::NTopic::TSessionClosedEvent>(&event)) { |
| 96 | + x->Confirm(); |
| 97 | + Cerr << ">>>>> " << x->DebugString() << Endl << Flush; |
| 98 | + } |
| 99 | + |
| 100 | + if (count == expected) { |
| 101 | + return; |
| 102 | + } |
| 103 | + } |
| 104 | + Sleep(TDuration::MilliSeconds(250)); |
| 105 | + } |
| 106 | + |
| 107 | + UNIT_ASSERT_VALUES_EQUAL(expected, count); |
| 108 | + } |
| 109 | + |
| 110 | + Y_UNIT_TEST(CDC_Write) { |
| 111 | + TTopicSdkTestSetup setup = CreateSetup(); |
| 112 | + auto client = setup.MakeClient(); |
| 113 | + auto tableClient = setup.MakeTableClient(); |
| 114 | + auto session = tableClient.CreateSession().GetValueSync().GetSession(); |
| 115 | + |
| 116 | + ExecuteQuery(session, R"( |
| 117 | + --!syntax_v1 |
| 118 | + CREATE TABLE `/Root/origin` ( |
| 119 | + id UInt64, |
| 120 | + order UInt64, |
| 121 | + value Text, |
| 122 | + PRIMARY KEY (id, order) |
| 123 | + ) WITH ( |
| 124 | + AUTO_PARTITIONING_MIN_PARTITIONS_COUNT = 64, |
| 125 | + AUTO_PARTITIONING_MAX_PARTITIONS_COUNT = 64, |
| 126 | + UNIFORM_PARTITIONS = 64 |
| 127 | + ); |
| 128 | + )"); |
| 129 | + |
| 130 | + ExecuteQuery(session, R"( |
| 131 | + --!syntax_v1 |
| 132 | + ALTER TABLE `/Root/origin` |
| 133 | + ADD CHANGEFEED `feed` WITH ( |
| 134 | + MODE = 'UPDATES', |
| 135 | + FORMAT = 'JSON', |
| 136 | + TOPIC_AUTO_PARTITIONING = 'ENABLED', |
| 137 | + TOPIC_MIN_ACTIVE_PARTITIONS = 2 |
| 138 | + ); |
| 139 | + )"); |
| 140 | + |
| 141 | + { |
| 142 | + TAlterTopicSettings alterSettings; |
| 143 | + alterSettings |
| 144 | + .BeginAlterPartitioningSettings() |
| 145 | + .MinActivePartitions(1) |
| 146 | + .MaxActivePartitions(10000) |
| 147 | + .BeginAlterAutoPartitioningSettings() |
| 148 | + .Strategy(EAutoPartitioningStrategy::ScaleUp) |
| 149 | + .StabilizationWindow(TDuration::Seconds(1)) |
| 150 | + .DownUtilizationPercent(1) |
| 151 | + .UpUtilizationPercent(2) |
| 152 | + .EndAlterAutoPartitioningSettings() |
| 153 | + .EndAlterTopicPartitioningSettings() |
| 154 | + .BeginAddConsumer() |
| 155 | + .ConsumerName("consumer-1") |
| 156 | + .EndAddConsumer(); |
| 157 | + auto f = client.AlterTopic("/Root/origin/feed", alterSettings); |
| 158 | + f.Wait(); |
| 159 | + |
| 160 | + auto v = f.GetValueSync(); |
| 161 | + UNIT_ASSERT_C(v.IsSuccess(), "Error: " << v); |
| 162 | + } |
| 163 | + |
| 164 | + Cerr << ">>>>> " << TInstant::Now() << " Start table insert" << Endl << Flush; |
| 165 | + ExecuteDataQuery(session, R"( |
| 166 | + --!syntax_v1 |
| 167 | + $sample = AsList( |
| 168 | + AsStruct(ListFromRange(0, 150000) AS v) |
| 169 | + ); |
| 170 | +
|
| 171 | + UPSERT INTO `/Root/origin` (id, order, value) |
| 172 | + SELECT |
| 173 | + RandomNumber(v) AS id, |
| 174 | + v AS order, |
| 175 | + CAST('0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF' AS Utf8?) AS value |
| 176 | + FROM as_table($sample) |
| 177 | + FLATTEN BY (v); |
| 178 | + )"); |
| 179 | + |
| 180 | + Cerr << ">>>>> " << TInstant::Now() << " Start read topic" << Endl << Flush; |
| 181 | + AssertMessageCountInTopic(client, "/Root/origin/feed/streamImpl", 150000); |
| 182 | + Cerr << ">>>>> " << TInstant::Now() << " End" << Endl << Flush; |
| 183 | + } |
| 184 | +} |
| 185 | + |
| 186 | +} // namespace NKikimr |
0 commit comments