Skip to content

Commit b3fc5c8

Browse files
GrigoriyPAGrigoriyPA
authored
Merge 4770fa9 into a0de26d 
2 parents a0de26d + 4770fa9 commit b3fc5c8

File tree

5 files changed

+103
-28
lines changed

5 files changed

+103
-28
lines changed

‎ydb/tests/tools/kqprun/kqprun.cpp‎

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -210,7 +210,7 @@ void RunScript(const TExecutionOptions& executionOptions, const NKqpRun::TRunner
210210
}
211211
}
212212

213-
if (runnerOptions.YdbSettings.MonitoringEnabled) {
213+
if (runnerOptions.YdbSettings.MonitoringEnabled || runnerOptions.YdbSettings.GrpcEnabled) {
214214
RunAsDaemon();
215215
}
216216

@@ -452,8 +452,17 @@ class TMain : public TMainClassArgs {
452452
});
453453
options.AddLongOption("inflight-limit", "In flight limit for async queries (use 0 for unlimited)")
454454
.RequiredArgument("uint")
455-
.DefaultValue(RunnerOptions.YdbSettings.InFlightLimit)
456-
.StoreResult(&RunnerOptions.YdbSettings.InFlightLimit);
455+
.DefaultValue(0)
456+
.StoreResult(&RunnerOptions.YdbSettings.AsyncQueriesSettings.InFlightLimit);
457+
TChoices<NKqpRun::TAsyncQueriesSettings::EVerbose> verbose({
458+
{"each-query", NKqpRun::TAsyncQueriesSettings::EVerbose::EachQuery},
459+
{"final", NKqpRun::TAsyncQueriesSettings::EVerbose::Final}
460+
});
461+
options.AddLongOption("async-verbose", "Verbose type for async queries")
462+
.RequiredArgument("type")
463+
.DefaultValue("each-query")
464+
.Choices(verbose.GetChoices())
465+
.StoreMappedResultT<TString>(&RunnerOptions.YdbSettings.AsyncQueriesSettings.Verbose, verbose);
457466

458467
TChoices<NKikimrKqp::EQueryAction> scriptAction({
459468
{"execute", NKikimrKqp::QUERY_ACTION_EXECUTE},
@@ -501,7 +510,7 @@ class TMain : public TMainClassArgs {
501510
return nodeCount;
502511
});
503512

504-
options.AddLongOption('M', "monitoring", "Embedded UI port (use 0 to start on random free port), if used kqprun will be runs as daemon")
513+
options.AddLongOption('M', "monitoring", "Embedded UI port (use 0 to start on random free port), if used kqprun will be run as daemon")
505514
.RequiredArgument("uint")
506515
.Handler1([this](const NLastGetopt::TOptsParser* option) {
507516
if (const TString& port = option->CurVal()) {
@@ -510,6 +519,15 @@ class TMain : public TMainClassArgs {
510519
}
511520
});
512521

522+
options.AddLongOption('G', "grpc", "gRPC port (use 0 to start on random free port), if used kqprun will be run as daemon")
523+
.RequiredArgument("uint")
524+
.Handler1([this](const NLastGetopt::TOptsParser* option) {
525+
if (const TString& port = option->CurVal()) {
526+
RunnerOptions.YdbSettings.GrpcEnabled = true;
527+
RunnerOptions.YdbSettings.GrpcPort = FromString(port);
528+
}
529+
});
530+
513531
options.AddLongOption('E', "emulate-yt", "Emulate YT tables (use file gateway instead of native gateway)")
514532
.NoArgument()
515533
.SetFlag(&EmulateYt);
@@ -529,7 +547,7 @@ class TMain : public TMainClassArgs {
529547
}
530548

531549
int DoRun(NLastGetopt::TOptsParseResult&&) override {
532-
if (!ExecutionOptions.SchemeQuery && ExecutionOptions.ScriptQueries.empty() && !RunnerOptions.YdbSettings.MonitoringEnabled) {
550+
if (!ExecutionOptions.SchemeQuery && ExecutionOptions.ScriptQueries.empty() && !RunnerOptions.YdbSettings.MonitoringEnabled && !RunnerOptions.YdbSettings.GrpcEnabled) {
533551
ythrow yexception() << "Nothing to execute";
534552
}
535553

‎ydb/tests/tools/kqprun/src/actors.cpp‎

Lines changed: 42 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -102,12 +102,17 @@ class TRunScriptActorMock : public NActors::TActorBootstrapped<TRunScriptActorMo
102102
class TAsyncQueryRunnerActor : public NActors::TActor<TAsyncQueryRunnerActor> {
103103
using TBase = NActors::TActor<TAsyncQueryRunnerActor>;
104104

105+
struct TRequestInfo {
106+
TInstant StartTime;
107+
NThreading::TFuture<TQueryResponse> RequestFuture;
108+
};
109+
105110
public:
106-
TAsyncQueryRunnerActor(ui64 inFlightLimit)
111+
TAsyncQueryRunnerActor(const TAsyncQueriesSettings& settings)
107112
: TBase(&TAsyncQueryRunnerActor::StateFunc)
108-
, InFlightLimit_(inFlightLimit)
113+
, Settings_(settings)
109114
{
110-
RunningRequests_.reserve(InFlightLimit_);
115+
RunningRequests_.reserve(Settings_.InFlightLimit);
111116
}
112117

113118
STRICT_STFUNC(StateFunc,
@@ -123,21 +128,29 @@ class TAsyncQueryRunnerActor : public NActors::TActor<TAsyncQueryRunnerActor> {
123128

124129
void Handle(TEvPrivate::TEvAsyncQueryFinished::TPtr& ev) {
125130
const ui64 requestId = ev->Get()->RequestId;
131+
RequestsLatency_ += TInstant::Now() - RunningRequests_[requestId].StartTime;
126132
RunningRequests_.erase(requestId);
127133

128134
const auto& response = ev->Get()->Result.Response->Get()->Record.GetRef();
129135
const auto status = response.GetYdbStatus();
130136

131137
if (status == Ydb::StatusIds::SUCCESS) {
132138
Completed_++;
133-
Cout << CoutColors_.Green() << TInstant::Now().ToIsoStringLocal() << " Request #" << requestId << " completed. " << CoutColors_.Yellow() << GetInfoString() << CoutColors_.Default() << Endl;
139+
if (Settings_.Verbose == TAsyncQueriesSettings::EVerbose::EachQuery) {
140+
Cout << CoutColors_.Green() << TInstant::Now().ToIsoStringLocal() << " Request #" << requestId << " completed. " << CoutColors_.Yellow() << GetInfoString() << CoutColors_.Default() << Endl;
141+
}
134142
} else {
135143
Failed_++;
136144
NYql::TIssues issues;
137145
NYql::IssuesFromMessage(response.GetResponse().GetQueryIssues(), issues);
138146
Cout << CoutColors_.Red() << TInstant::Now().ToIsoStringLocal() << " Request #" << requestId << " failed " << status << ". " << CoutColors_.Yellow() << GetInfoString() << "\n" << CoutColors_.Red() << "Issues:\n" << issues.ToString() << CoutColors_.Default();
139147
}
140148

149+
if (Settings_.Verbose == TAsyncQueriesSettings::EVerbose::Final && TInstant::Now() - LastReportTime_ > TDuration::Seconds(1)) {
150+
Cout << CoutColors_.Green() << TInstant::Now().ToIsoStringLocal() << " Finished " << Failed_ + Completed_ << " requests. " << CoutColors_.Yellow() << GetInfoString() << CoutColors_.Default() << Endl;
151+
LastReportTime_ = TInstant::Now();
152+
}
153+
141154
StartDelayedRequests();
142155
TryFinalize();
143156
}
@@ -151,18 +164,23 @@ class TAsyncQueryRunnerActor : public NActors::TActor<TAsyncQueryRunnerActor> {
151164

152165
private:
153166
void StartDelayedRequests() {
154-
while (!DelayedRequests_.empty() && (!InFlightLimit_ || RunningRequests_.size() < InFlightLimit_)) {
167+
while (!DelayedRequests_.empty() && (!Settings_.InFlightLimit || RunningRequests_.size() < Settings_.InFlightLimit)) {
155168
auto request = std::move(DelayedRequests_.front());
156169
DelayedRequests_.pop();
157170

158171
auto promise = NThreading::NewPromise<TQueryResponse>();
159172
Register(CreateRunScriptActorMock(std::move(request->Get()->Request), promise, nullptr));
160-
RunningRequests_[RequestId_] = promise.GetFuture().Subscribe([id = RequestId_, this](const NThreading::TFuture<TQueryResponse>& f) {
161-
Send(SelfId(), new TEvPrivate::TEvAsyncQueryFinished(id, std::move(f.GetValue())));
162-
});
173+
RunningRequests_[RequestId_] = {
174+
.StartTime = TInstant::Now(),
175+
.RequestFuture = promise.GetFuture().Subscribe([id = RequestId_, this](const NThreading::TFuture<TQueryResponse>& f) {
176+
Send(SelfId(), new TEvPrivate::TEvAsyncQueryFinished(id, std::move(f.GetValue())));
177+
})
178+
};
163179

164180
MaxInFlight_ = std::max(MaxInFlight_, RunningRequests_.size());
165-
Cout << TStringBuilder() << CoutColors_.Cyan() << TInstant::Now().ToIsoStringLocal() << " Request #" << RequestId_ << " started. " << CoutColors_.Yellow() << GetInfoString() << CoutColors_.Default() << "\n";
181+
if (Settings_.Verbose == TAsyncQueriesSettings::EVerbose::EachQuery) {
182+
Cout << TStringBuilder() << CoutColors_.Cyan() << TInstant::Now().ToIsoStringLocal() << " Request #" << RequestId_ << " started. " << CoutColors_.Yellow() << GetInfoString() << CoutColors_.Default() << "\n";
183+
}
166184

167185
RequestId_++;
168186
request->Get()->StartPromise.SetValue();
@@ -174,28 +192,38 @@ class TAsyncQueryRunnerActor : public NActors::TActor<TAsyncQueryRunnerActor> {
174192
return false;
175193
}
176194

195+
if (Settings_.Verbose == TAsyncQueriesSettings::EVerbose::Final) {
196+
Cout << TStringBuilder() << CoutColors_.Cyan() << TInstant::Now().ToIsoStringLocal() << " All async requests finished. " << CoutColors_.Yellow() << GetInfoString() << CoutColors_.Default() << "\n";
197+
}
198+
177199
FinalizePromise_->SetValue();
178200
PassAway();
179201
return true;
180202
}
181203

182204
TString GetInfoString() const {
183-
return TStringBuilder() << "completed: " << Completed_ << ", failed: " << Failed_ << ", in flight: " << RunningRequests_.size() << ", max in flight: " << MaxInFlight_ << ", spend time: " << TInstant::Now() - StartTime_;
205+
TStringBuilder result = TStringBuilder() << "completed: " << Completed_ << ", failed: " << Failed_ << ", in flight: " << RunningRequests_.size() << ", max in flight: " << MaxInFlight_ << ", spend time: " << TInstant::Now() - StartTime_;
206+
if (const auto amountRequests = Completed_ + Failed_) {
207+
result << ", average latency: " << RequestsLatency_ / amountRequests;
208+
}
209+
return result;
184210
}
185211

186212
private:
187-
const ui64 InFlightLimit_;
213+
const TAsyncQueriesSettings Settings_;
188214
const TInstant StartTime_ = TInstant::Now();
189215
const NColorizer::TColors CoutColors_ = NColorizer::AutoColors(Cout);
190216

191217
std::optional<NThreading::TPromise<void>> FinalizePromise_;
192218
std::queue<TEvPrivate::TEvStartAsyncQuery::TPtr> DelayedRequests_;
193-
std::unordered_map<ui64, NThreading::TFuture<TQueryResponse>> RunningRequests_;
219+
std::unordered_map<ui64, TRequestInfo> RunningRequests_;
220+
TInstant LastReportTime_ = TInstant::Now();
194221

195222
ui64 RequestId_ = 1;
196223
ui64 MaxInFlight_ = 0;
197224
ui64 Completed_ = 0;
198225
ui64 Failed_ = 0;
226+
TDuration RequestsLatency_;
199227
};
200228

201229
class TResourcesWaiterActor : public NActors::TActorBootstrapped<TResourcesWaiterActor> {
@@ -270,8 +298,8 @@ NActors::IActor* CreateRunScriptActorMock(TQueryRequest request, NThreading::TPr
270298
return new TRunScriptActorMock(std::move(request), promise, progressCallback);
271299
}
272300

273-
NActors::IActor* CreateAsyncQueryRunnerActor(ui64 inFlightLimit) {
274-
return new TAsyncQueryRunnerActor(inFlightLimit);
301+
NActors::IActor* CreateAsyncQueryRunnerActor(const TAsyncQueriesSettings& settings) {
302+
return new TAsyncQueryRunnerActor(settings);
275303
}
276304

277305
NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount) {

‎ydb/tests/tools/kqprun/src/actors.h‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#pragma once
22

3+
#include "common.h"
4+
35
#include <ydb/core/kqp/common/events/events.h>
46
#include <ydb/core/kqp/executer_actor/kqp_executer.h>
57

@@ -72,7 +74,7 @@ using TProgressCallback = std::function<void(const NKikimrKqp::TEvExecuterProgre
7274

7375
NActors::IActor* CreateRunScriptActorMock(TQueryRequest request, NThreading::TPromise<TQueryResponse> promise, TProgressCallback progressCallback);
7476

75-
NActors::IActor* CreateAsyncQueryRunnerActor(ui64 inFlightLimit);
77+
NActors::IActor* CreateAsyncQueryRunnerActor(const TAsyncQueriesSettings& settings);
7678

7779
NActors::IActor* CreateResourcesWaiterActor(NThreading::TPromise<void> promise, i32 expectedNodeCount);
7880

‎ydb/tests/tools/kqprun/src/common.h‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,27 @@ namespace NKqpRun {
1414

1515
constexpr char YQL_TOKEN_VARIABLE[] = "YQL_TOKEN";
1616

17+
struct TAsyncQueriesSettings {
18+
enum class EVerbose {
19+
EachQuery,
20+
Final,
21+
};
22+
23+
ui64 InFlightLimit = 0;
24+
EVerbose Verbose = EVerbose::EachQuery;
25+
};
26+
1727
struct TYdbSetupSettings {
1828
ui32 NodeCount = 1;
1929
TString DomainName = "Root";
2030
TDuration InitializationTimeout = TDuration::Seconds(10);
2131

2232
bool MonitoringEnabled = false;
2333
ui16 MonitoringPortOffset = 0;
34+
35+
bool GrpcEnabled = false;
36+
ui16 GrpcPort = 0;
37+
2438
bool TraceOptEnabled = false;
2539
TString LogOutputFile;
2640

@@ -29,8 +43,7 @@ struct TYdbSetupSettings {
2943
NKikimr::NMiniKQL::TComputationNodeFactory ComputationFactory;
3044
TIntrusivePtr<NYql::IYtGateway> YtGateway;
3145
NKikimrConfig::TAppConfig AppConfig;
32-
33-
ui64 InFlightLimit = 0;
46+
TAsyncQueriesSettings AsyncQueriesSettings;
3447
};
3548

3649

‎ydb/tests/tools/kqprun/src/ydb_setup.cpp‎

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -120,8 +120,8 @@ class TYdbSetup::TImpl {
120120
serverSettings.SetFrFactory(functionRegistryFactory);
121121
}
122122

123-
NKikimr::Tests::TServerSettings GetServerSettings() {
124-
ui32 msgBusPort = PortManager_.GetPort();
123+
NKikimr::Tests::TServerSettings GetServerSettings(ui32 grpcPort) {
124+
const ui32 msgBusPort = PortManager_.GetPort();
125125

126126
NKikimr::Tests::TServerSettings serverSettings(msgBusPort, Settings_.AppConfig.GetAuthConfig(), Settings_.AppConfig.GetPQConfig());
127127
serverSettings.SetNodeCount(Settings_.NodeCount);
@@ -152,15 +152,23 @@ class TYdbSetup::TImpl {
152152
serverSettings.SetNeedStatsCollectors(true);
153153
}
154154

155+
if (Settings_.GrpcEnabled) {
156+
serverSettings.SetGrpcPort(grpcPort);
157+
}
158+
155159
return serverSettings;
156160
}
157161

158-
void InitializeServer() {
159-
NKikimr::Tests::TServerSettings serverSettings = GetServerSettings();
162+
void InitializeServer(ui32 grpcPort) {
163+
NKikimr::Tests::TServerSettings serverSettings = GetServerSettings(grpcPort);
160164

161165
Server_ = MakeHolder<NKikimr::Tests::TServer>(serverSettings);
162166
Server_->GetRuntime()->SetDispatchTimeout(TDuration::Max());
163167

168+
if (Settings_.GrpcEnabled) {
169+
Server_->EnableGRpc(grpcPort);
170+
}
171+
164172
Client_ = MakeHolder<NKikimr::Tests::TClient>(serverSettings);
165173
Client_->InitRootScheme();
166174
}
@@ -204,15 +212,21 @@ class TYdbSetup::TImpl {
204212
: Settings_(settings)
205213
, CoutColors_(NColorizer::AutoColors(Cout))
206214
{
215+
const ui32 grpcPort = Settings_.GrpcPort ? Settings_.GrpcPort : PortManager_.GetPort();
216+
207217
InitializeYqlLogger();
208-
InitializeServer();
218+
InitializeServer(grpcPort);
209219
WaitResourcesPublishing();
210220

211221
if (Settings_.MonitoringEnabled) {
212222
for (ui32 nodeIndex = 0; nodeIndex < Settings_.NodeCount; ++nodeIndex) {
213223
Cout << CoutColors_.Cyan() << "Monitoring port" << (Settings_.NodeCount > 1 ? TStringBuilder() << " for node " << nodeIndex + 1 : TString()) << ": " << CoutColors_.Default() << Server_->GetRuntime()->GetMonPort(nodeIndex) << Endl;
214224
}
215225
}
226+
227+
if (Settings_.GrpcEnabled) {
228+
Cout << CoutColors_.Cyan() << "gRPC port: " << CoutColors_.Default() << grpcPort << Endl;
229+
}
216230
}
217231

218232
NKikimr::NKqp::TEvKqp::TEvQueryResponse::TPtr SchemeQueryRequest(const TRequestOptions& query) const {
@@ -274,7 +288,7 @@ class TYdbSetup::TImpl {
274288

275289
void QueryRequestAsync(const TRequestOptions& query) {
276290
if (!AsyncQueryRunnerActorId_) {
277-
AsyncQueryRunnerActorId_ = GetRuntime()->Register(CreateAsyncQueryRunnerActor(Settings_.InFlightLimit));
291+
AsyncQueryRunnerActorId_ = GetRuntime()->Register(CreateAsyncQueryRunnerActor(Settings_.AsyncQueriesSettings));
278292
}
279293

280294
auto request = GetQueryRequest(query);

0 commit comments

Comments
 (0)