Skip to content

fq/checkpoint storage: enable more query caching #18453

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 33 additions & 28 deletions ydb/core/fq/libs/checkpoint_storage/ydb_checkpoint_storage.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -59,18 +59,15 @@ struct TCheckpointContext : public TThrRefBase {
TGenerationContextPtr GenerationContext;
TCheckpointGraphDescriptionContextPtr CheckpointGraphDescriptionContext;
IEntityIdGenerator::TPtr EntityIdGenerator;
TExecDataQuerySettings Settings;

TCheckpointContext(const TCheckpointId& id,
ECheckpointStatus status,
ECheckpointStatus expected,
ui64 stateSizeBytes,
TExecDataQuerySettings settings)
ui64 stateSizeBytes)
: CheckpointId(id)
, Status(status)
, ExpectedStatus(expected)
, StateSizeBytes(stateSizeBytes)
, Settings(settings)
{
}
};
Expand Down Expand Up @@ -115,7 +112,8 @@ TFuture<TDataQueryResult> SelectGraphCoordinators(const TGenerationContextPtr& c

return context->Session.ExecuteDataQuery(
query,
TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx());
TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(),
context->ExecDataQuerySettings);
}

TFuture<TStatus> ProcessCoordinators(
Expand Down Expand Up @@ -223,7 +221,7 @@ TFuture<TStatus> CreateCheckpoint(const TCheckpointContextPtr& context) {
}

auto ttxControl = TTxControl::Tx(*generationContext->Transaction).CommitTx();
return generationContext->Session.ExecuteDataQuery(query, ttxControl, params.Build(), context->Settings).Apply(
return generationContext->Session.ExecuteDataQuery(query, ttxControl, params.Build(), generationContext->ExecDataQuerySettings).Apply(
[] (const TFuture<TDataQueryResult>& future) {
TStatus status = future.GetValue();
return status;
Expand Down Expand Up @@ -273,7 +271,7 @@ TFuture<TStatus> UpdateCheckpoint(const TCheckpointContextPtr& context) {
.Build();

auto ttxControl = TTxControl::Tx(*generationContext->Transaction).CommitTx();
return generationContext->Session.ExecuteDataQuery(query, ttxControl, params.Build(), context->Settings).Apply(
return generationContext->Session.ExecuteDataQuery(query, ttxControl, params.Build(), generationContext->ExecDataQuerySettings).Apply(
[] (const TFuture<TDataQueryResult>& future) {
TStatus status = future.GetValue();
return status;
Expand All @@ -300,7 +298,7 @@ TFuture<TDataQueryResult> SelectGraphDescId(const TCheckpointContextPtr& context
.String(graphDescContext->GraphDescId)
.Build();

return generationContext->Session.ExecuteDataQuery(query, TTxControl::Tx(*generationContext->Transaction), params.Build(), context->Settings);
return generationContext->Session.ExecuteDataQuery(query, TTxControl::Tx(*generationContext->Transaction), params.Build(), generationContext->ExecDataQuerySettings);
}

bool GraphDescIdExists(const TFuture<TDataQueryResult>& result) {
Expand Down Expand Up @@ -352,7 +350,7 @@ TFuture<TStatus> CreateCheckpointWrapper(
});
}

TFuture<TDataQueryResult> SelectGraphCheckpoints(const TGenerationContextPtr& context, const TVector<ECheckpointStatus>& statuses, ui64 limit, TExecDataQuerySettings settings, bool loadGraphDescription)
TFuture<TDataQueryResult> SelectGraphCheckpoints(const TGenerationContextPtr& context, const TVector<ECheckpointStatus>& statuses, ui64 limit, bool loadGraphDescription)
{
NYdb::TParamsBuilder paramsBuilder;
if (statuses) {
Expand Down Expand Up @@ -418,7 +416,7 @@ TFuture<TDataQueryResult> SelectGraphCheckpoints(const TGenerationContextPtr& co
query,
TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx(),
params,
settings);
context->ExecDataQuerySettings);
}

TFuture<TStatus> ProcessCheckpoints(
Expand Down Expand Up @@ -498,7 +496,7 @@ TFuture<TDataQueryResult> SelectCheckpoint(const TCheckpointContextPtr& context)
query,
TTxControl::Tx(*generationContext->Transaction),
params.Build(),
context->Settings);
generationContext->ExecDataQuerySettings);
}

TFuture<TStatus> CheckCheckpoint(
Expand Down Expand Up @@ -756,7 +754,8 @@ TFuture<TIssues> TCheckpointStorage::Init()
TFuture<TIssues> TCheckpointStorage::RegisterGraphCoordinator(const TCoordinatorId& coordinator)
{
auto future = YdbConnection->TableClient.RetryOperation(
[prefix = YdbConnection->TablePathPrefix, coordinator] (TSession session) {
[prefix = YdbConnection->TablePathPrefix, coordinator,
execDataQuerySettings = DefaultExecDataQuerySettings()] (TSession session) {
auto context = MakeIntrusive<TGenerationContext>(
session,
true,
Expand All @@ -765,7 +764,8 @@ TFuture<TIssues> TCheckpointStorage::RegisterGraphCoordinator(const TCoordinator
"graph_id",
"generation",
coordinator.GraphId,
coordinator.Generation);
coordinator.Generation,
execDataQuerySettings);

return RegisterCheckGeneration(context);
});
Expand All @@ -777,7 +777,7 @@ TFuture<ICheckpointStorage::TGetCoordinatorsResult> TCheckpointStorage::GetCoord
auto getContext = MakeIntrusive<TGetCoordinatorsContext>();

auto future = YdbConnection->TableClient.RetryOperation(
[prefix = YdbConnection->TablePathPrefix, getContext] (TSession session) {
[prefix = YdbConnection->TablePathPrefix, getContext, execDataQuerySettings = DefaultExecDataQuerySettings()] (TSession session) {
auto generationContext = MakeIntrusive<TGenerationContext>(
session,
false,
Expand All @@ -786,7 +786,8 @@ TFuture<ICheckpointStorage::TGetCoordinatorsResult> TCheckpointStorage::GetCoord
"graph_id",
"generation",
"",
0UL);
0UL,
execDataQuerySettings);

auto future = SelectGraphCoordinators(generationContext);
return future.Apply(
Expand All @@ -811,7 +812,7 @@ TFuture<ICheckpointStorage::TCreateCheckpointResult> TCheckpointStorage::CreateC
ECheckpointStatus status)
{
Y_ABORT_UNLESS(graphDescId);
auto checkpointContext = MakeIntrusive<TCheckpointContext>(checkpointId, status, ECheckpointStatus::Pending, 0ul, DefaultExecDataQuerySettings());
auto checkpointContext = MakeIntrusive<TCheckpointContext>(checkpointId, status, ECheckpointStatus::Pending, 0ul);
checkpointContext->CheckpointGraphDescriptionContext = MakeIntrusive<TCheckpointGraphDescriptionContext>(graphDescId);
return CreateCheckpointImpl(coordinator, checkpointContext);
}
Expand All @@ -822,7 +823,7 @@ TFuture<ICheckpointStorage::TCreateCheckpointResult> TCheckpointStorage::CreateC
const NProto::TCheckpointGraphDescription& graphDesc,
ECheckpointStatus status)
{
auto checkpointContext = MakeIntrusive<TCheckpointContext>(checkpointId, status, ECheckpointStatus::Pending, 0ul, DefaultExecDataQuerySettings());
auto checkpointContext = MakeIntrusive<TCheckpointContext>(checkpointId, status, ECheckpointStatus::Pending, 0ul);
checkpointContext->CheckpointGraphDescriptionContext = MakeIntrusive<TCheckpointGraphDescriptionContext>(graphDesc);
checkpointContext->EntityIdGenerator = EntityIdGenerator;
return CreateCheckpointImpl(coordinator, checkpointContext);
Expand All @@ -831,7 +832,7 @@ TFuture<ICheckpointStorage::TCreateCheckpointResult> TCheckpointStorage::CreateC
TFuture<ICheckpointStorage::TCreateCheckpointResult> TCheckpointStorage::CreateCheckpointImpl(const TCoordinatorId& coordinator, const TCheckpointContextPtr& checkpointContext) {
Y_ABORT_UNLESS(checkpointContext->CheckpointGraphDescriptionContext->GraphDescId || checkpointContext->EntityIdGenerator);
auto future = YdbConnection->TableClient.RetryOperation(
[prefix = YdbConnection->TablePathPrefix, coordinator, checkpointContext] (TSession session) {
[prefix = YdbConnection->TablePathPrefix, coordinator, checkpointContext, execDataQuerySettings = DefaultExecDataQuerySettings()] (TSession session) {
auto generationContext = MakeIntrusive<TGenerationContext>(
session,
false,
Expand All @@ -840,7 +841,8 @@ TFuture<ICheckpointStorage::TCreateCheckpointResult> TCheckpointStorage::CreateC
"graph_id",
"generation",
coordinator.GraphId,
coordinator.Generation);
coordinator.Generation,
execDataQuerySettings);

checkpointContext->GenerationContext = generationContext;

Expand All @@ -863,9 +865,9 @@ TFuture<TIssues> TCheckpointStorage::UpdateCheckpointStatus(
ECheckpointStatus prevStatus,
ui64 stateSizeBytes)
{
auto checkpointContext = MakeIntrusive<TCheckpointContext>(checkpointId, newStatus, prevStatus, stateSizeBytes, DefaultExecDataQuerySettings());
auto checkpointContext = MakeIntrusive<TCheckpointContext>(checkpointId, newStatus, prevStatus, stateSizeBytes);
auto future = YdbConnection->TableClient.RetryOperation(
[prefix = YdbConnection->TablePathPrefix, coordinator, checkpointContext] (TSession session) {
[prefix = YdbConnection->TablePathPrefix, coordinator, checkpointContext, execDataQuerySettings = DefaultExecDataQuerySettings()] (TSession session) {
auto generationContext = MakeIntrusive<TGenerationContext>(
session,
false,
Expand All @@ -874,7 +876,8 @@ TFuture<TIssues> TCheckpointStorage::UpdateCheckpointStatus(
"graph_id",
"generation",
coordinator.GraphId,
coordinator.Generation);
coordinator.Generation,
execDataQuerySettings);

checkpointContext->GenerationContext = generationContext;

Expand All @@ -889,9 +892,9 @@ TFuture<TIssues> TCheckpointStorage::AbortCheckpoint(
const TCoordinatorId& coordinator,
const TCheckpointId& checkpointId)
{
auto checkpointContext = MakeIntrusive<TCheckpointContext>(checkpointId, ECheckpointStatus::Aborted, ECheckpointStatus::Pending, 0ul, DefaultExecDataQuerySettings());
auto checkpointContext = MakeIntrusive<TCheckpointContext>(checkpointId, ECheckpointStatus::Aborted, ECheckpointStatus::Pending, 0ul);
auto future = YdbConnection->TableClient.RetryOperation(
[prefix = YdbConnection->TablePathPrefix, coordinator, checkpointContext] (TSession session) {
[prefix = YdbConnection->TablePathPrefix, coordinator, checkpointContext, execDataQuerySettings = DefaultExecDataQuerySettings()] (TSession session) {
auto generationContext = MakeIntrusive<TGenerationContext>(
session,
false,
Expand All @@ -900,7 +903,8 @@ TFuture<TIssues> TCheckpointStorage::AbortCheckpoint(
"graph_id",
"generation",
coordinator.GraphId,
coordinator.Generation);
coordinator.Generation,
execDataQuerySettings);

checkpointContext->GenerationContext = generationContext;

Expand All @@ -921,7 +925,7 @@ TFuture<ICheckpointStorage::TGetCheckpointsResult> TCheckpointStorage::GetCheckp
auto getContext = MakeIntrusive<TGetCheckpointsContext>();

auto future = YdbConnection->TableClient.RetryOperation(
[prefix = YdbConnection->TablePathPrefix, graph, getContext, statuses, limit, loadGraphDescription, settings = DefaultExecDataQuerySettings()] (TSession session) {
[prefix = YdbConnection->TablePathPrefix, graph, getContext, statuses, limit, loadGraphDescription, execDataQuerySettings = DefaultExecDataQuerySettings()] (TSession session) {
auto generationContext = MakeIntrusive<TGenerationContext>(
session,
false,
Expand All @@ -930,9 +934,10 @@ TFuture<ICheckpointStorage::TGetCheckpointsResult> TCheckpointStorage::GetCheckp
"graph_id",
"generation",
graph,
0UL);
0UL,
execDataQuerySettings);

auto future = SelectGraphCheckpoints(generationContext, statuses, limit, settings, loadGraphDescription);
auto future = SelectGraphCheckpoints(generationContext, statuses, limit, loadGraphDescription);
return future.Apply(
[generationContext, getContext, loadGraphDescription] (const TFuture<TDataQueryResult>& future) {
return ProcessCheckpoints(future.GetValue(), generationContext, getContext, loadGraphDescription);
Expand Down
33 changes: 24 additions & 9 deletions ydb/core/fq/libs/ydb/ydb.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -25,23 +25,29 @@ TFuture<TDataQueryResult> SelectGeneration(const TGenerationContextPtr& context)
auto query = Sprintf(R"(
--!syntax_v1
PRAGMA TablePathPrefix("%s");
DECLARE $pk AS String;

SELECT %s, %s
FROM %s
WHERE %s = "%s";
WHERE %s = $pk;
)", context->TablePathPrefix.c_str(),
context->PrimaryKeyColumn.c_str(),
context->GenerationColumn.c_str(),
context->Table.c_str(),
context->PrimaryKeyColumn.c_str(),
context->PrimaryKey.c_str());
context->PrimaryKeyColumn.c_str());

NYdb::TParamsBuilder params;
params
.AddParam("$pk")
.String(context->PrimaryKey)
.Build();

auto ttxControl = TTxControl::BeginTx(TTxSettings::SerializableRW());
if (context->OperationType == TGenerationContext::Check && context->CommitTx) {
ttxControl.CommitTx();
}

return context->Session.ExecuteDataQuery(query, ttxControl);
return context->Session.ExecuteDataQuery(query, ttxControl, params.Build(), context->ExecDataQuerySettings);
}

TFuture<TStatus> CheckGeneration(
Expand Down Expand Up @@ -127,23 +133,32 @@ TFuture<TStatus> UpsertGeneration(const TGenerationContextPtr& context) {
auto query = Sprintf(R"(
--!syntax_v1
PRAGMA TablePathPrefix("%s");
DECLARE $pk AS String;
DECLARE $generation AS Uint64;

UPSERT INTO %s (%s, %s) VALUES
("%s", %lu);
($pk, $generation);
)", context->TablePathPrefix.c_str(),
context->Table.c_str(),
context->PrimaryKeyColumn.c_str(),
context->GenerationColumn.c_str(),
context->PrimaryKey.c_str(),
context->Generation);
context->GenerationColumn.c_str());

NYdb::TParamsBuilder params;
params
.AddParam("$pk")
.String(context->PrimaryKey)
.Build()
.AddParam("$generation")
.Uint64(context->Generation)
.Build();

auto ttxControl = TTxControl::Tx(*context->Transaction);
if (context->CommitTx) {
ttxControl.CommitTx();
context->Transaction.reset();
}

return context->Session.ExecuteDataQuery(query, ttxControl).Apply(
return context->Session.ExecuteDataQuery(query, ttxControl, params.Build(), context->ExecDataQuerySettings).Apply(
[] (const TFuture<TDataQueryResult>& future) {
TStatus status = future.GetValue();
return status;
Expand Down
6 changes: 5 additions & 1 deletion ydb/core/fq/libs/ydb/ydb.h
Original file line number Diff line number Diff line change
Expand Up @@ -80,14 +80,17 @@ struct TGenerationContext : public TThrRefBase {
// result of Select
ui64 GenerationRead = 0;

NYdb::NTable::TExecDataQuerySettings ExecDataQuerySettings;

TGenerationContext(NYdb::NTable::TSession session,
bool commitTx,
const TString& tablePathPrefix,
const TString& table,
const TString& primaryKeyColumn,
const TString& generationColumn,
const TString& primaryKey,
ui64 generation)
ui64 generation,
const NYdb::NTable::TExecDataQuerySettings& execDataQuerySettings = {})
: Session(session)
, CommitTx(commitTx)
, TablePathPrefix(tablePathPrefix)
Expand All @@ -96,6 +99,7 @@ struct TGenerationContext : public TThrRefBase {
, GenerationColumn(generationColumn)
, PrimaryKey(primaryKey)
, Generation(generation)
, ExecDataQuerySettings(execDataQuerySettings)
{
}
};
Expand Down
Loading