Merge 6a917f9 into 1bb608a

nikvas0 · web-flow · commit b07f192a8761 · 2024-09-06T17:01:28.000+03:00
diff --git a/ydb/core/kqp/common/kqp_tx.cpp b/ydb/core/kqp/common/kqp_tx.cpp
@@ -7,24 +7,33 @@ namespace NKqp {
 
 using namespace NYql;
 
-TIssue GetLocksInvalidatedIssue(const TKqpTransactionContext& txCtx, const TMaybe<TKqpTxLock>& invalidatedLock) {
+NYql::TIssue GetLocksInvalidatedIssue(const TKqpTransactionContext& txCtx, const TKikimrPathId& pathId) {
     TStringBuilder message;
     message << "Transaction locks invalidated.";
 
-    TMaybe<TString> tableName;
-    if (invalidatedLock) {
-        TKikimrPathId id(invalidatedLock->GetSchemeShard(), invalidatedLock->GetPathId());
-        auto table = txCtx.TableByIdMap.FindPtr(id);
-        if (table) {
-            tableName = *table;
+    if (pathId.OwnerId() != 0) {
+        auto table = txCtx.TableByIdMap.FindPtr(pathId);
+        if (!table) {
+            return YqlIssue(TPosition(), TIssuesIds::KIKIMR_LOCKS_INVALIDATED, message << " Unknown table.");    
         }
+        return YqlIssue(TPosition(), TIssuesIds::KIKIMR_LOCKS_INVALIDATED, message << " Table: " << *table);
+    } else {
+        // Olap tables don't return SchemeShard in locks, thus we use tableId here.
+        for (const auto& [pathId, table] : txCtx.TableByIdMap) {
+            if (pathId.TableId() == pathId.TableId()) {
+                return YqlIssue(TPosition(), TIssuesIds::KIKIMR_LOCKS_INVALIDATED, message << " Table: " << table);
+            }
+        }
+        return YqlIssue(TPosition(), TIssuesIds::KIKIMR_LOCKS_INVALIDATED, message << " Unknown table.");    
     }
+}
 
-    if (tableName) {
-        message << " Table: " << *tableName;
-    }
-
-    return YqlIssue(TPosition(), TIssuesIds::KIKIMR_LOCKS_INVALIDATED, message);
+TIssue GetLocksInvalidatedIssue(const TKqpTransactionContext& txCtx, const TKqpTxLock& invalidatedLock) {
+    return GetLocksInvalidatedIssue(
+        txCtx,
+        TKikimrPathId(
+            invalidatedLock.GetSchemeShard(),
+            invalidatedLock.GetPathId()));
 }
 
 std::pair<bool, std::vector<TIssue>> MergeLocks(const NKikimrMiniKQL::TType& type, const NKikimrMiniKQL::TValue& value,
diff --git a/ydb/core/kqp/common/kqp_tx.h b/ydb/core/kqp/common/kqp_tx.h
@@ -434,6 +434,7 @@ class TTransactionsCache {
     }
 };
 
+NYql::TIssue GetLocksInvalidatedIssue(const TKqpTransactionContext& txCtx, const NYql::TKikimrPathId& pathId);
 std::pair<bool, std::vector<NYql::TIssue>> MergeLocks(const NKikimrMiniKQL::TType& type,
     const NKikimrMiniKQL::TValue& value, TKqpTransactionContext& txCtx);
 
diff --git a/ydb/core/kqp/executer_actor/kqp_data_executer.cpp b/ydb/core/kqp/executer_actor/kqp_data_executer.cpp
@@ -241,11 +241,11 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
     }
 
     void Finalize() {
+        YQL_ENSURE(!AlreadyReplied);
         if (LocksBroken) {
-            TString message = "Transaction locks invalidated.";
-
-            return ReplyErrorAndDie(Ydb::StatusIds::ABORTED,
-                YqlIssue({}, TIssuesIds::KIKIMR_LOCKS_INVALIDATED, message));
+            return ReplyErrorAndDie(
+                Ydb::StatusIds::ABORTED,
+                YqlIssue(TPosition(), TIssuesIds::KIKIMR_LOCKS_INVALIDATED, "Transaction locks invalidated. Unknown table."));
         }
 
         auto& response = *ResponseEv->Record.MutableResponse();
@@ -1183,6 +1183,23 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                 CheckExecutionComplete();
                 return;
             }
+            case NKikimrDataEvents::TEvWriteResult::STATUS_LOCKS_BROKEN: {
+                LOG_D("Broken locks: " << res->Record.DebugString());
+                YQL_ENSURE(shardState->State == TShardState::EState::Executing);
+                shardState->State = TShardState::EState::Finished;
+                Counters->TxProxyMon->TxResultAborted->Inc();
+                LocksBroken = true;
+
+                if (!res->Record.GetTxLocks().empty()) {
+                    ResponseEv->BrokenLockPathId = NYql::TKikimrPathId(
+                        res->Record.GetTxLocks(0).GetSchemeShard(),
+                        res->Record.GetTxLocks(0).GetPathId());
+                    return ReplyErrorAndDie(Ydb::StatusIds::ABORTED, {});
+                }
+
+                CheckExecutionComplete();
+                return;
+            }
             default:
             {
                 return ShardError(res->Record);
@@ -1235,29 +1252,15 @@ class TKqpDataExecuter : public TKqpExecuterBase<TKqpDataExecuter, EExecType::Da
                 shardState->State = TShardState::EState::Finished;
 
                 Counters->TxProxyMon->TxResultAborted->Inc(); // TODO: dedicated counter?
-
                 LocksBroken = true;
 
-                TMaybe<TString> tableName;
                 if (!res->Record.GetTxLocks().empty()) {
-                    auto& lock = res->Record.GetTxLocks(0);
-                    auto tableId = TTableId(lock.GetSchemeShard(), lock.GetPathId());
-                    auto it = FindIf(TasksGraph.GetStagesInfo(), [tableId](const auto& x){ return x.second.Meta.TableId.HasSamePath(tableId); });
-                    if (it != TasksGraph.GetStagesInfo().end()) {
-                        tableName = it->second.Meta.TableConstInfo->Path;
-                    }
-                }
-
-                // Reply as soon as we know which table had locks invalidated
-                if (tableName) {
-                    auto message = TStringBuilder()
-                        << "Transaction locks invalidated. Table: " << *tableName;
-
-                    return ReplyErrorAndDie(Ydb::StatusIds::ABORTED,
-                        YqlIssue({}, TIssuesIds::KIKIMR_LOCKS_INVALIDATED, message));
+                    ResponseEv->BrokenLockPathId = TKikimrPathId(
+                        res->Record.GetTxLocks(0).GetSchemeShard(),
+                        res->Record.GetTxLocks(0).GetPathId());
+                    return ReplyErrorAndDie(Ydb::StatusIds::ABORTED, {});
                 }
 
-                // Receive more replies from other shards
                 CheckExecutionComplete();
                 return;
             }
diff --git a/ydb/core/kqp/executer_actor/kqp_executer.h b/ydb/core/kqp/executer_actor/kqp_executer.h
@@ -26,6 +26,7 @@ struct TEvKqpExecuter {
 
         NLWTrace::TOrbit Orbit;
         IKqpGateway::TKqpSnapshot Snapshot;
+        std::optional<NYql::TKikimrPathId> BrokenLockPathId;
         ui64 ResultRowsCount = 0;
         ui64 ResultRowsBytes = 0;
 
diff --git a/ydb/core/kqp/executer_actor/kqp_executer_impl.h b/ydb/core/kqp/executer_actor/kqp_executer_impl.h
@@ -1768,7 +1768,9 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
         auto& response = *ResponseEv->Record.MutableResponse();
 
         response.SetStatus(status);
-        response.MutableIssues()->Swap(issues);
+        if (issues) {
+            response.MutableIssues()->Swap(issues);
+        }
 
         LOG_T("ReplyErrorAndDie. Response: " << response.DebugString()
             << ", to ActorId: " << Target);
diff --git a/ydb/core/kqp/provider/yql_kikimr_provider.h b/ydb/core/kqp/provider/yql_kikimr_provider.h
@@ -323,7 +323,6 @@ class TKikimrTransactionContextBase : public TThrRefBase {
 
         for (const auto& info : tableInfos) {
             tableInfoMap.emplace(info.GetTableName(), &info);
-
             TKikimrPathId pathId(info.GetTableId().GetOwnerId(), info.GetTableId().GetTableId());
             TableByIdMap.emplace(pathId, info.GetTableName());
         }
diff --git a/ydb/core/kqp/session_actor/kqp_session_actor.cpp b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
@@ -1439,6 +1439,12 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
 
             // Invalidate query cache on scheme/internal errors
             switch (status) {
+                case Ydb::StatusIds::ABORTED: {
+                    if (ev->BrokenLockPathId) {
+                        issues.AddIssue(GetLocksInvalidatedIssue(*QueryState->TxCtx, *ev->BrokenLockPathId));
+                    }
+                    break;
+                }
                 case Ydb::StatusIds::SCHEME_ERROR:
                 case Ydb::StatusIds::INTERNAL_ERROR:
                     InvalidateQuery();
diff --git a/ydb/core/kqp/ut/tx/kqp_locks_ut.cpp b/ydb/core/kqp/ut/tx/kqp_locks_ut.cpp
@@ -204,6 +204,42 @@ Y_UNIT_TEST_SUITE(KqpLocks) {
         UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
         CompareYson(R"([[[2u];#;[11u];["Session2"]]])", FormatResultSetYson(result.GetResultSet(0)));
     }
+
+    Y_UNIT_TEST(TwoPhaseTx) {
+        TKikimrRunner kikimr;
+        auto db = kikimr.GetTableClient();
+
+        auto session1 = db.CreateSession().GetValueSync().GetSession();
+        auto session2 = db.CreateSession().GetValueSync().GetSession();
+
+        auto result = session1.ExecuteDataQuery(Q_(R"(
+            REPLACE INTO `/Root/Test` (Group, Name, Comment) VALUES (1U, "Paul", "Changed");
+            SELECT * FROM `/Root/Test` WHERE Name == "Paul" ORDER BY Group, Name;
+        )"), TTxControl::BeginTx(TTxSettings::SerializableRW())).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        auto tx1 = result.GetTransaction();
+        UNIT_ASSERT(tx1);
+
+        result = session2.ExecuteDataQuery(Q_(R"(
+            REPLACE INTO `/Root/Test` (Group, Name, Comment)
+            VALUES (1U, "Paul", "Changed");
+        )"), TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx()).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        result = session1.ExecuteDataQuery(Q_(R"(
+            SELECT * FROM `KeyValue`;
+        )"), TTxControl::Tx(*tx1)).ExtractValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
+
+        auto commitResult = tx1->Commit().GetValueSync();
+        UNIT_ASSERT_VALUES_EQUAL_C(commitResult.GetStatus(), EStatus::ABORTED, result.GetIssues().ToString());
+        commitResult.GetIssues().PrintTo(Cerr);
+        UNIT_ASSERT_C(HasIssue(commitResult.GetIssues(), NYql::TIssuesIds::KIKIMR_LOCKS_INVALIDATED,
+            [] (const NYql::TIssue& issue) {
+                return issue.GetMessage().Contains("/Root/Test");
+            }), commitResult.GetIssues().ToString());
+    }
 }
 
 } // namespace NKqp

Original file line number	Diff line number	Diff line change
`@@ -434,6 +434,7 @@ class TTransactionsCache {`
`434`	`434`	`}`
`435`	`435`	`};`
`436`	`436`
	`437`	`+NYql::TIssue GetLocksInvalidatedIssue(const TKqpTransactionContext& txCtx, const NYql::TKikimrPathId& pathId);`
`437`	`438`	`std::pair<bool, std::vector<NYql::TIssue>> MergeLocks(const NKikimrMiniKQL::TType& type,`
`438`	`439`	`const NKikimrMiniKQL::TValue& value, TKqpTransactionContext& txCtx);`
`439`	`440`
Original file line number	Diff line number	Diff line change
`@@ -323,7 +323,6 @@ class TKikimrTransactionContextBase : public TThrRefBase {`
`323`	`323`
`324`	`324`	`for (const auto& info : tableInfos) {`
`325`	`325`	`tableInfoMap.emplace(info.GetTableName(), &info);`
`326`		`-`
`327`	`326`	`TKikimrPathId pathId(info.GetTableId().GetOwnerId(), info.GetTableId().GetTableId());`
`328`	`327`	`TableByIdMap.emplace(pathId, info.GetTableName());`
`329`	`328`	`}`