Skip to content

Commit 9d12242

Browse files
GrigoriyPAGrigoriyPA
authored
YQ WM fix for sls feature flag (#7057)
1 parent a87900c commit 9d12242

File tree

6 files changed

+220
-38
lines changed

6 files changed

+220
-38
lines changed

‎ydb/core/kqp/workload_service/actors/actors.h

+5-2
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,15 @@ namespace NKikimr::NKqp::NWorkload {
99
NActors::IActor* CreatePoolHandlerActor(const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, NMonitoring::TDynamicCounterPtr counters);
1010

1111
// Fetch pool and create default pool if needed
12-
NActors::IActor* CreatePoolResolverActor(TEvPlaceRequestIntoPool::TPtr event, bool defaultPoolExists, bool enableOnServerless);
12+
NActors::IActor* CreatePoolResolverActor(TEvPlaceRequestIntoPool::TPtr event, bool defaultPoolExists);
1313

1414
// Fetch and create pool in scheme shard
15-
NActors::IActor* CreatePoolFetcherActor(const NActors::TActorId& replyActorId, const TString& database, const TString& poolId, TIntrusiveConstPtr<NACLib::TUserToken> userToken, bool enableOnServerless);
15+
NActors::IActor* CreatePoolFetcherActor(const NActors::TActorId& replyActorId, const TString& database, const TString& poolId, TIntrusiveConstPtr<NACLib::TUserToken> userToken);
1616
NActors::IActor* CreatePoolCreatorActor(const NActors::TActorId& replyActorId, const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, TIntrusiveConstPtr<NACLib::TUserToken> userToken, NACLibProto::TDiffACL diffAcl);
1717

18+
// Checks that database is serverless
19+
NActors::IActor* CreateDatabaseFetcherActor(const NActors::TActorId& replyActorId, const TString& database);
20+
1821
// Cpu load fetcher actor
1922
NActors::IActor* CreateCpuLoadFetcherActor(const NActors::TActorId& replyActorId);
2023

‎ydb/core/kqp/workload_service/actors/scheme_actors.cpp

+102-16
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,8 @@ using namespace NActors;
2222

2323
class TPoolResolverActor : public TActorBootstrapped<TPoolResolverActor> {
2424
public:
25-
TPoolResolverActor(TEvPlaceRequestIntoPool::TPtr event, bool defaultPoolExists, bool enableOnServerless)
25+
TPoolResolverActor(TEvPlaceRequestIntoPool::TPtr event, bool defaultPoolExists)
2626
: Event(std::move(event))
27-
, EnableOnServerless(enableOnServerless)
2827
{
2928
if (!Event->Get()->PoolId) {
3029
Event->Get()->PoolId = NResourcePool::DEFAULT_POOL_ID;
@@ -39,7 +38,7 @@ class TPoolResolverActor : public TActorBootstrapped<TPoolResolverActor> {
3938

4039
void StartPoolFetchRequest() const {
4140
LOG_D("Start pool fetching");
42-
Register(CreatePoolFetcherActor(SelfId(), Event->Get()->Database, Event->Get()->PoolId, Event->Get()->UserToken, EnableOnServerless));
41+
Register(CreatePoolFetcherActor(SelfId(), Event->Get()->Database, Event->Get()->PoolId, Event->Get()->UserToken));
4342
}
4443

4544
void Handle(TEvPrivate::TEvFetchPoolResponse::TPtr& ev) {
@@ -116,20 +115,18 @@ class TPoolResolverActor : public TActorBootstrapped<TPoolResolverActor> {
116115

117116
private:
118117
TEvPlaceRequestIntoPool::TPtr Event;
119-
const bool EnableOnServerless;
120118
bool CanCreatePool = false;
121119
bool DefaultPoolCreated = false;
122120
};
123121

124122

125123
class TPoolFetcherActor : public TSchemeActorBase<TPoolFetcherActor> {
126124
public:
127-
TPoolFetcherActor(const TActorId& replyActorId, const TString& database, const TString& poolId, TIntrusiveConstPtr<NACLib::TUserToken> userToken, bool enableOnServerless)
125+
TPoolFetcherActor(const TActorId& replyActorId, const TString& database, const TString& poolId, TIntrusiveConstPtr<NACLib::TUserToken> userToken)
128126
: ReplyActorId(replyActorId)
129127
, Database(database)
130128
, PoolId(poolId)
131129
, UserToken(userToken)
132-
, EnableOnServerless(enableOnServerless)
133130
{}
134131

135132
void DoBootstrap() {
@@ -144,11 +141,6 @@ class TPoolFetcherActor : public TSchemeActorBase<TPoolFetcherActor> {
144141
}
145142

146143
const auto& result = results[0];
147-
if (!EnableOnServerless && result.DomainInfo && result.DomainInfo->IsServerless()) {
148-
Reply(Ydb::StatusIds::UNSUPPORTED, "Resource pools are disabled for serverless domains. Please contact your system administrator to enable it");
149-
return;
150-
}
151-
152144
switch (result.Status) {
153145
case EStatus::Unknown:
154146
case EStatus::PathNotTable:
@@ -238,7 +230,6 @@ class TPoolFetcherActor : public TSchemeActorBase<TPoolFetcherActor> {
238230
const TString Database;
239231
const TString PoolId;
240232
const TIntrusiveConstPtr<NACLib::TUserToken> UserToken;
241-
const bool EnableOnServerless;
242233

243234
NResourcePool::TPoolSettings PoolConfig;
244235
NKikimrProto::TPathID PathId;
@@ -451,18 +442,113 @@ class TPoolCreatorActor : public TSchemeActorBase<TPoolCreatorActor> {
451442
TActorId SchemePipeActorId;
452443
};
453444

445+
446+
class TDatabaseFetcherActor : public TSchemeActorBase<TDatabaseFetcherActor> {
447+
public:
448+
TDatabaseFetcherActor(const TActorId& replyActorId, const TString& database)
449+
: ReplyActorId(replyActorId)
450+
, Database(database)
451+
{}
452+
453+
void DoBootstrap() {
454+
Become(&TDatabaseFetcherActor::StateFunc);
455+
}
456+
457+
void Handle(TEvTxProxySchemeCache::TEvNavigateKeySetResult::TPtr& ev) {
458+
const auto& results = ev->Get()->Request->ResultSet;
459+
if (results.size() != 1) {
460+
Reply(Ydb::StatusIds::INTERNAL_ERROR, "Unexpected scheme cache response");
461+
return;
462+
}
463+
464+
const auto& result = results[0];
465+
switch (result.Status) {
466+
case EStatus::Unknown:
467+
case EStatus::PathNotTable:
468+
case EStatus::PathNotPath:
469+
case EStatus::RedirectLookupError:
470+
case EStatus::AccessDenied:
471+
case EStatus::RootUnknown:
472+
case EStatus::PathErrorUnknown:
473+
Reply(Ydb::StatusIds::NOT_FOUND, TStringBuilder() << "Database " << Database << " not found or you don't have access permissions");
474+
return;
475+
case EStatus::LookupError:
476+
case EStatus::TableCreationNotComplete:
477+
if (!ScheduleRetry(TStringBuilder() << "Retry error " << result.Status)) {
478+
Reply(Ydb::StatusIds::UNAVAILABLE, TStringBuilder() << "Retry limit exceeded on scheme error: " << result.Status);
479+
}
480+
return;
481+
case EStatus::Ok:
482+
Serverless = result.DomainInfo && result.DomainInfo->IsServerless();
483+
Reply(Ydb::StatusIds::SUCCESS);
484+
return;
485+
}
486+
}
487+
488+
STFUNC(StateFunc) {
489+
switch (ev->GetTypeRewrite()) {
490+
hFunc(TEvTxProxySchemeCache::TEvNavigateKeySetResult, Handle);
491+
default:
492+
StateFuncBase(ev);
493+
}
494+
}
495+
496+
protected:
497+
void StartRequest() override {
498+
LOG_D("Start database fetching");
499+
auto event = NTableCreator::BuildSchemeCacheNavigateRequest({{}}, Database, nullptr);
500+
event->ResultSet[0].Operation = NSchemeCache::TSchemeCacheNavigate::OpPath;
501+
Send(MakeSchemeCacheID(), new TEvTxProxySchemeCache::TEvNavigateKeySet(event.Release()), IEventHandle::FlagTrackDelivery);
502+
}
503+
504+
void OnFatalError(Ydb::StatusIds::StatusCode status, NYql::TIssue issue) override {
505+
Reply(status, {std::move(issue)});
506+
}
507+
508+
TString LogPrefix() const override {
509+
return TStringBuilder() << "[TDatabaseFetcherActor] ActorId: " << SelfId() << ", Database: " << Database << ", ";
510+
}
511+
512+
private:
513+
void Reply(Ydb::StatusIds::StatusCode status, const TString& message) {
514+
Reply(status, {NYql::TIssue(message)});
515+
}
516+
517+
void Reply(Ydb::StatusIds::StatusCode status, NYql::TIssues issues = {}) {
518+
if (status == Ydb::StatusIds::SUCCESS) {
519+
LOG_D("Database info successfully fetched");
520+
} else {
521+
LOG_W("Failed to fetch database info, " << status << ", issues: " << issues.ToOneLineString());
522+
}
523+
524+
Issues.AddIssues(std::move(issues));
525+
Send(ReplyActorId, new TEvPrivate::TEvFetchDatabaseResponse(status, Database, Serverless, std::move(Issues)));
526+
PassAway();
527+
}
528+
529+
private:
530+
const TActorId ReplyActorId;
531+
const TString Database;
532+
533+
bool Serverless = false;
534+
};
535+
454536
} // anonymous namespace
455537

456-
IActor* CreatePoolResolverActor(TEvPlaceRequestIntoPool::TPtr event, bool defaultPoolExists, bool enableOnServerless) {
457-
return new TPoolResolverActor(std::move(event), defaultPoolExists, enableOnServerless);
538+
IActor* CreatePoolResolverActor(TEvPlaceRequestIntoPool::TPtr event, bool defaultPoolExists) {
539+
return new TPoolResolverActor(std::move(event), defaultPoolExists);
458540
}
459541

460-
IActor* CreatePoolFetcherActor(const TActorId& replyActorId, const TString& database, const TString& poolId, TIntrusiveConstPtr<NACLib::TUserToken> userToken, bool enableOnServerless) {
461-
return new TPoolFetcherActor(replyActorId, database, poolId, userToken, enableOnServerless);
542+
IActor* CreatePoolFetcherActor(const TActorId& replyActorId, const TString& database, const TString& poolId, TIntrusiveConstPtr<NACLib::TUserToken> userToken) {
543+
return new TPoolFetcherActor(replyActorId, database, poolId, userToken);
462544
}
463545

464546
IActor* CreatePoolCreatorActor(const TActorId& replyActorId, const TString& database, const TString& poolId, const NResourcePool::TPoolSettings& poolConfig, TIntrusiveConstPtr<NACLib::TUserToken> userToken, NACLibProto::TDiffACL diffAcl) {
465547
return new TPoolCreatorActor(replyActorId, database, poolId, poolConfig, userToken, diffAcl);
466548
}
467549

550+
IActor* CreateDatabaseFetcherActor(const TActorId& replyActorId, const TString& database) {
551+
return new TDatabaseFetcherActor(replyActorId, database);
552+
}
553+
468554
} // NKikimr::NKqp::NWorkload

‎ydb/core/kqp/workload_service/common/events.h

+15
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ struct TEvPrivate {
2222
EvRefreshPoolState = EventSpaceBegin(NActors::TEvents::ES_PRIVATE),
2323
EvResolvePoolResponse,
2424
EvFetchPoolResponse,
25+
EvFetchDatabaseResponse,
2526
EvCreatePoolResponse,
2627
EvPrepareTablesRequest,
2728
EvPlaceRequestIntoPoolResponse,
@@ -85,6 +86,20 @@ struct TEvPrivate {
8586
const NYql::TIssues Issues;
8687
};
8788

89+
struct TEvFetchDatabaseResponse : public NActors::TEventLocal<TEvFetchDatabaseResponse, EvFetchDatabaseResponse> {
90+
TEvFetchDatabaseResponse(Ydb::StatusIds::StatusCode status, const TString& database, bool serverless, NYql::TIssues issues)
91+
: Status(status)
92+
, Database(database)
93+
, Serverless(serverless)
94+
, Issues(std::move(issues))
95+
{}
96+
97+
const Ydb::StatusIds::StatusCode Status;
98+
const TString Database;
99+
const bool Serverless;
100+
const NYql::TIssues Issues;
101+
};
102+
88103
struct TEvCreatePoolResponse : public NActors::TEventLocal<TEvCreatePoolResponse, EvCreatePoolResponse> {
89104
TEvCreatePoolResponse(Ydb::StatusIds::StatusCode status, NYql::TIssues issues)
90105
: Status(status)

‎ydb/core/kqp/workload_service/kqp_workload_service.cpp

+41-19
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,23 @@ using namespace NActors;
2727

2828

2929
class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
30+
struct TCounters {
31+
const NMonitoring::TDynamicCounterPtr Counters;
32+
33+
NMonitoring::TDynamicCounters::TCounterPtr ActivePools;
34+
35+
TCounters(NMonitoring::TDynamicCounterPtr counters)
36+
: Counters(counters)
37+
{
38+
Register();
39+
}
40+
41+
private:
42+
void Register() {
43+
ActivePools = Counters->GetCounter("ActivePools", false);
44+
}
45+
};
46+
3047
enum class ETablesCreationStatus {
3148
Cleanup,
3249
NotStarted,
@@ -43,9 +60,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
4360
public:
4461
explicit TKqpWorkloadService(NMonitoring::TDynamicCounterPtr counters)
4562
: Counters(counters)
46-
{
47-
RegisterCounters();
48-
}
63+
{}
4964

5065
void Bootstrap() {
5166
Become(&TKqpWorkloadService::MainState);
@@ -55,7 +70,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
5570
(ui32)NKikimrConsole::TConfigItem::FeatureFlagsItem
5671
}), IEventHandle::FlagTrackDelivery);
5772

58-
CpuQuotaManager = std::make_unique<TCpuQuotaManagerState>(ActorContext(), Counters->GetSubgroup("subcomponent", "CpuQuotaManager"));
73+
CpuQuotaManager = std::make_unique<TCpuQuotaManagerState>(ActorContext(), Counters.Counters->GetSubgroup("subcomponent", "CpuQuotaManager"));
5974

6075
EnabledResourcePools = AppData()->FeatureFlags.GetEnableResourcePools();
6176
EnabledResourcePoolsOnServerless = AppData()->FeatureFlags.GetEnableResourcePoolsOnServerless();
@@ -132,9 +147,9 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
132147
return;
133148
}
134149

135-
LOG_D("Recieved new request from " << workerActorId << ", Database: " << ev->Get()->Database << ", PoolId: " << ev->Get()->PoolId << ", SessionId: " << ev->Get()->SessionId);
136-
bool hasDefaultPool = DatabasesWithDefaultPool.contains(CanonizePath(ev->Get()->Database));
137-
Register(CreatePoolResolverActor(std::move(ev), hasDefaultPool, EnabledResourcePoolsOnServerless));
150+
const TString& database = ev->Get()->Database;
151+
LOG_D("Recieved new request from " << workerActorId << ", Database: " << database << ", PoolId: " << ev->Get()->PoolId << ", SessionId: " << ev->Get()->SessionId);
152+
GetOrCreateDatabaseState(database)->DoPlaceRequest(std::move(ev));
138153
}
139154

140155
void Handle(TEvCleanupRequest::TPtr& ev) {
@@ -177,6 +192,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
177192
hFunc(TEvCleanupRequest, Handle);
178193
hFunc(TEvents::TEvWakeup, Handle);
179194

195+
hFunc(TEvPrivate::TEvFetchDatabaseResponse, Handle);
180196
hFunc(TEvPrivate::TEvResolvePoolResponse, Handle);
181197
hFunc(TEvPrivate::TEvPlaceRequestIntoPoolResponse, Handle);
182198
hFunc(TEvPrivate::TEvNodesInfoRequest, Handle);
@@ -191,11 +207,15 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
191207
)
192208

193209
private:
210+
void Handle(TEvPrivate::TEvFetchDatabaseResponse::TPtr& ev) {
211+
GetOrCreateDatabaseState(ev->Get()->Database)->UpdateDatabaseInfo(ev);
212+
}
213+
194214
void Handle(TEvPrivate::TEvResolvePoolResponse::TPtr& ev) {
195215
const auto& event = ev->Get()->Event;
196216
const TString& database = event->Get()->Database;
197217
if (ev->Get()->DefaultPoolCreated) {
198-
DatabasesWithDefaultPool.insert(CanonizePath(database));
218+
GetOrCreateDatabaseState(database)->HasDefaultPool = true;
199219
}
200220

201221
const TString& poolId = event->Get()->PoolId;
@@ -211,10 +231,10 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
211231
TString poolKey = GetPoolKey(database, poolId);
212232
LOG_I("Creating new handler for pool " << poolKey);
213233

214-
auto poolHandler = Register(CreatePoolHandlerActor(database, poolId, ev->Get()->PoolConfig, Counters));
234+
auto poolHandler = Register(CreatePoolHandlerActor(database, poolId, ev->Get()->PoolConfig, Counters.Counters));
215235
poolState = &PoolIdToState.insert({poolKey, TPoolState{.PoolHandler = poolHandler, .ActorContext = ActorContext()}}).first->second;
216236

217-
ActivePools->Inc();
237+
Counters.ActivePools->Inc();
218238
ScheduleIdleCheck();
219239
}
220240

@@ -409,7 +429,7 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
409429
}
410430
for (const auto& poolKey : poolsToDelete) {
411431
PoolIdToState.erase(poolKey);
412-
ActivePools->Dec();
432+
Counters.ActivePools->Dec();
413433
}
414434

415435
if (!PoolIdToState.empty()) {
@@ -472,6 +492,14 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
472492
Send(replyActorId, new TEvCleanupResponse(status, {NYql::TIssue(message)}));
473493
}
474494

495+
TDatabaseState* GetOrCreateDatabaseState(const TString& database) {
496+
auto databaseIt = DatabaseToState.find(database);
497+
if (databaseIt != DatabaseToState.end()) {
498+
return &databaseIt->second;
499+
}
500+
return &DatabaseToState.insert({database, TDatabaseState{.ActorContext = ActorContext(), .EnabledResourcePoolsOnServerless = EnabledResourcePoolsOnServerless}}).first->second;
501+
}
502+
475503
TPoolState* GetPoolState(const TString& database, const TString& poolId) {
476504
return GetPoolState(GetPoolKey(database, poolId));
477505
}
@@ -492,12 +520,8 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
492520
return "[Service] ";
493521
}
494522

495-
void RegisterCounters() {
496-
ActivePools = Counters->GetCounter("ActivePools", false);
497-
}
498-
499523
private:
500-
NMonitoring::TDynamicCounterPtr Counters;
524+
TCounters Counters;
501525

502526
bool EnabledResourcePools = false;
503527
bool EnabledResourcePoolsOnServerless = false;
@@ -506,12 +530,10 @@ class TKqpWorkloadService : public TActorBootstrapped<TKqpWorkloadService> {
506530
ETablesCreationStatus TablesCreationStatus = ETablesCreationStatus::Cleanup;
507531
std::unordered_set<TString> PendingHandlers;
508532

509-
std::unordered_set<TString> DatabasesWithDefaultPool;
533+
std::unordered_map<TString, TDatabaseState> DatabaseToState;
510534
std::unordered_map<TString, TPoolState> PoolIdToState;
511535
std::unique_ptr<TCpuQuotaManagerState> CpuQuotaManager;
512536
ui32 NodeCount = 0;
513-
514-
NMonitoring::TDynamicCounters::TCounterPtr ActivePools;
515537
};
516538

517539
} // anonymous namespace

0 commit comments

Comments
 (0)