@@ -200,77 +200,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
200200 return ReplyErrorAndDie (Ydb::StatusIds::ABORTED, {});
201201 }
202202
203- auto addLocks = [this ](const ui64 taskId, const auto & data) {
204- if (data.GetData ().template Is <NKikimrTxDataShard::TEvKqpInputActorResultInfo>()) {
205- NKikimrTxDataShard::TEvKqpInputActorResultInfo info;
206- YQL_ENSURE (data.GetData ().UnpackTo (&info), " Failed to unpack settings" );
207- NDataIntegrity::LogIntegrityTrails (" InputActorResult" , Request.UserTraceId , TxId, info, TlsActivationContext->AsActorContext ());
208- for (auto & lock : info.GetLocks ()) {
209- if (!TxManager) {
210- Locks.push_back (lock);
211- }
212-
213- const auto & task = TasksGraph.GetTask (taskId);
214- const auto & stageInfo = TasksGraph.GetStageInfo (task.StageId );
215- ShardIdToTableInfo->Add (lock.GetDataShard (), stageInfo.Meta .TableKind == ETableKind::Olap, stageInfo.Meta .TablePath );
216-
217- if (TxManager) {
218- TxManager->AddShard (lock.GetDataShard (), stageInfo.Meta .TableKind == ETableKind::Olap, stageInfo.Meta .TablePath );
219- TxManager->AddAction (lock.GetDataShard (), IKqpTransactionManager::EAction::READ);
220- TxManager->AddLock (lock.GetDataShard (), lock);
221- }
222- }
223-
224- if (!BatchOperationSettings.Empty () && info.HasBatchOperationMaxKey ()) {
225- if (ResponseEv->BatchOperationMaxKeys .empty ()) {
226- for (auto keyId : info.GetBatchOperationKeyIds ()) {
227- ResponseEv->BatchOperationKeyIds .push_back (keyId);
228- }
229- }
230-
231- ResponseEv->BatchOperationMaxKeys .emplace_back (info.GetBatchOperationMaxKey ());
232- }
233- } else if (data.GetData ().template Is <NKikimrKqp::TEvKqpOutputActorResultInfo>()) {
234- NKikimrKqp::TEvKqpOutputActorResultInfo info;
235- YQL_ENSURE (data.GetData ().UnpackTo (&info), " Failed to unpack settings" );
236- NDataIntegrity::LogIntegrityTrails (" OutputActorResult" , Request.UserTraceId , TxId, info, TlsActivationContext->AsActorContext ());
237- for (auto & lock : info.GetLocks ()) {
238- if (!TxManager) {
239- Locks.push_back (lock);
240- }
241-
242- const auto & task = TasksGraph.GetTask (taskId);
243- const auto & stageInfo = TasksGraph.GetStageInfo (task.StageId );
244- ShardIdToTableInfo->Add (lock.GetDataShard (), stageInfo.Meta .TableKind == ETableKind::Olap, stageInfo.Meta .TablePath );
245- if (TxManager) {
246- YQL_ENSURE (stageInfo.Meta .TableKind == ETableKind::Olap);
247- IKqpTransactionManager::TActionFlags flags = IKqpTransactionManager::EAction::WRITE;
248- if (info.GetHasRead ()) {
249- flags |= IKqpTransactionManager::EAction::READ;
250- }
251-
252- TxManager->AddShard (lock.GetDataShard (), stageInfo.Meta .TableKind == ETableKind::Olap, stageInfo.Meta .TablePath );
253- TxManager->AddAction (lock.GetDataShard (), flags);
254- TxManager->AddLock (lock.GetDataShard (), lock);
255- }
256- }
257- }
258- };
259-
260- for (auto & [_, extraData] : ExtraData) {
261- for (const auto & source : extraData.Data .GetSourcesExtraData ()) {
262- addLocks (extraData.TaskId , source);
263- }
264- for (const auto & transform : extraData.Data .GetInputTransformsData ()) {
265- addLocks (extraData.TaskId , transform);
266- }
267- for (const auto & sink : extraData.Data .GetSinksExtraData ()) {
268- addLocks (extraData.TaskId , sink);
269- }
270- if (extraData.Data .HasComputeExtraData ()) {
271- addLocks (extraData.TaskId , extraData.Data .GetComputeExtraData ());
272- }
273- }
203+ FillLocksFromExtraData ();
274204
275205 if (TxManager) {
276206 TxManager->SetHasSnapshot (GetSnapshot ().IsValid ());
@@ -306,7 +236,6 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
306236 IEventHandle::FlagTrackDelivery,
307237 0 ,
308238 ExecuterSpan.GetTraceId ());
309- MakeResponseAndPassAway ();
310239 return ;
311240 } else if (Request.UseImmediateEffects ) {
312241 Become (&TKqpDataExecuter::FinalizeState);
@@ -378,9 +307,9 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
378307 MakeResponseAndPassAway ();
379308 }
380309
381- void HandleFinalize (TEvents::TEvUndelivered::TPtr&) {
382- auto issue = YqlIssue ({}, TIssuesIds::KIKIMR_TEMPORARILY_UNAVAILABLE, " Buffer actor isn't available. " );
383- ReplyErrorAndDie (Ydb::StatusIds::UNAVAILABLE, issue );
310+ void HandleFinalize (TEvents::TEvUndelivered::TPtr& ev ) {
311+ AFL_ENSURE (ev-> Sender == BufferActorId );
312+ LOG_W ( " Got Undelivered from BufferActor: " << ev-> Sender );
384313 }
385314
386315 void MakeResponseAndPassAway () {
@@ -1273,6 +1202,9 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
12731202
12741203 void Handle (TEvKqpBuffer::TEvError::TPtr& ev) {
12751204 auto & msg = *ev->Get ();
1205+ if (msg.Stats && Stats) {
1206+ Stats->AddBufferStats (std::move (*msg.Stats ));
1207+ }
12761208 TBase::HandleAbortExecution (msg.StatusCode , msg.Issues , false );
12771209 }
12781210
@@ -2820,6 +2752,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
28202752 }
28212753
28222754 if (Planner->GetPendingComputeTasks ().empty () && Planner->GetPendingComputeActors ().empty ()) {
2755+ FillLocksFromExtraData ();
28232756 PassAway ();
28242757 }
28252758 }
@@ -2829,6 +2762,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
28292762 LOG_I (" Timed out on waiting for Compute Actors to finish - forcing shutdown. Sender: " << ev->Sender );
28302763
28312764 if (ev->Sender == SelfId ()) {
2765+ FillLocksFromExtraData ();
28322766 PassAway ();
28332767 }
28342768 }
@@ -2840,6 +2774,7 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
28402774
28412775 // In case of external timeout the response is already sent to the client - no need to wait for stats.
28422776 if (statusCode == Ydb::StatusIds::TIMEOUT) {
2777+ FillLocksFromExtraData ();
28432778 LOG_I (" External timeout while waiting for Compute Actors to finish - forcing shutdown. Sender: " << ev->Sender );
28442779 PassAway ();
28452780 }
@@ -2944,6 +2879,80 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
29442879 }
29452880 }
29462881
2882+ void FillLocksFromExtraData () {
2883+ auto addLocks = [this ](const ui64 taskId, const auto & data) {
2884+ if (data.GetData ().template Is <NKikimrTxDataShard::TEvKqpInputActorResultInfo>()) {
2885+ NKikimrTxDataShard::TEvKqpInputActorResultInfo info;
2886+ YQL_ENSURE (data.GetData ().UnpackTo (&info), " Failed to unpack settings" );
2887+ NDataIntegrity::LogIntegrityTrails (" InputActorResult" , Request.UserTraceId , TxId, info, TlsActivationContext->AsActorContext ());
2888+ for (auto & lock : info.GetLocks ()) {
2889+ if (!TxManager) {
2890+ Locks.push_back (lock);
2891+ }
2892+
2893+ const auto & task = TasksGraph.GetTask (taskId);
2894+ const auto & stageInfo = TasksGraph.GetStageInfo (task.StageId );
2895+ ShardIdToTableInfo->Add (lock.GetDataShard (), stageInfo.Meta .TableKind == ETableKind::Olap, stageInfo.Meta .TablePath );
2896+
2897+ if (TxManager) {
2898+ TxManager->AddShard (lock.GetDataShard (), stageInfo.Meta .TableKind == ETableKind::Olap, stageInfo.Meta .TablePath );
2899+ TxManager->AddAction (lock.GetDataShard (), IKqpTransactionManager::EAction::READ);
2900+ TxManager->AddLock (lock.GetDataShard (), lock);
2901+ }
2902+ }
2903+
2904+ if (!BatchOperationSettings.Empty () && info.HasBatchOperationMaxKey ()) {
2905+ if (ResponseEv->BatchOperationMaxKeys .empty ()) {
2906+ for (auto keyId : info.GetBatchOperationKeyIds ()) {
2907+ ResponseEv->BatchOperationKeyIds .push_back (keyId);
2908+ }
2909+ }
2910+
2911+ ResponseEv->BatchOperationMaxKeys .emplace_back (info.GetBatchOperationMaxKey ());
2912+ }
2913+ } else if (data.GetData ().template Is <NKikimrKqp::TEvKqpOutputActorResultInfo>()) {
2914+ NKikimrKqp::TEvKqpOutputActorResultInfo info;
2915+ YQL_ENSURE (data.GetData ().UnpackTo (&info), " Failed to unpack settings" );
2916+ NDataIntegrity::LogIntegrityTrails (" OutputActorResult" , Request.UserTraceId , TxId, info, TlsActivationContext->AsActorContext ());
2917+ for (auto & lock : info.GetLocks ()) {
2918+ if (!TxManager) {
2919+ Locks.push_back (lock);
2920+ }
2921+
2922+ const auto & task = TasksGraph.GetTask (taskId);
2923+ const auto & stageInfo = TasksGraph.GetStageInfo (task.StageId );
2924+ ShardIdToTableInfo->Add (lock.GetDataShard (), stageInfo.Meta .TableKind == ETableKind::Olap, stageInfo.Meta .TablePath );
2925+ if (TxManager) {
2926+ YQL_ENSURE (stageInfo.Meta .TableKind == ETableKind::Olap);
2927+ IKqpTransactionManager::TActionFlags flags = IKqpTransactionManager::EAction::WRITE;
2928+ if (info.GetHasRead ()) {
2929+ flags |= IKqpTransactionManager::EAction::READ;
2930+ }
2931+
2932+ TxManager->AddShard (lock.GetDataShard (), stageInfo.Meta .TableKind == ETableKind::Olap, stageInfo.Meta .TablePath );
2933+ TxManager->AddAction (lock.GetDataShard (), flags);
2934+ TxManager->AddLock (lock.GetDataShard (), lock);
2935+ }
2936+ }
2937+ }
2938+ };
2939+
2940+ for (auto & [_, extraData] : ExtraData) {
2941+ for (const auto & source : extraData.Data .GetSourcesExtraData ()) {
2942+ addLocks (extraData.TaskId , source);
2943+ }
2944+ for (const auto & transform : extraData.Data .GetInputTransformsData ()) {
2945+ addLocks (extraData.TaskId , transform);
2946+ }
2947+ for (const auto & sink : extraData.Data .GetSinksExtraData ()) {
2948+ addLocks (extraData.TaskId , sink);
2949+ }
2950+ if (extraData.Data .HasComputeExtraData ()) {
2951+ addLocks (extraData.TaskId , extraData.Data .GetComputeExtraData ());
2952+ }
2953+ }
2954+ }
2955+
29472956private:
29482957 TShardIdToTableInfoPtr ShardIdToTableInfo;
29492958
0 commit comments