[+] cache for SchemeNavigate

Author: Alek5andr-Kotov

ydb/core/kqp/session_actor/kqp_session_actor.cpp

@@ -240,6 +240,9 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
             << "ActorId: " << SelfId() << ", "
             << "ActorState: " << CurrentStateFuncName() << ", ";
         if (Y_LIKELY(QueryState)) {
+            if (QueryState->HasTxControl()) {
+                result << " TxId: " << QueryState->GetTxControl().tx_id() << ", ";
+            }
             result << "TraceId: " << QueryState->UserRequestContext->TraceId << ", ";
         }
         return result;
@@ -547,6 +550,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
     }
 
     void AddOffsetsToTransaction() {
+        LOG_I("begin request for TopicOperations");
         YQL_ENSURE(QueryState);
         if (!PrepareQueryTransaction()) {
             return;
@@ -557,6 +561,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
         if (!AreAllTheTopicsAndPartitionsKnown()) {
             auto navigate = QueryState->BuildSchemeCacheNavigate();
             Become(&TKqpSessionActor::ExecuteState);
+            LOG_I("begin request for SchemeNavigate");
             Send(MakeSchemeCacheID(), new TEvTxProxySchemeCache::TEvNavigateKeySet(navigate.release()));
             return;
         }
@@ -572,7 +577,10 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
             return;
         }
 
+        LOG_I("end request for TopicOperations");
         ReplySuccess();
+
+        LOG_I("after ReplySuccess");
     }
 
     void CompileQuery() {
@@ -2909,8 +2917,10 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
     }
 
     void ProcessTopicOps(TEvTxProxySchemeCache::TEvNavigateKeySetResult::TPtr& ev) {
+        LOG_I("end request for SchemeNavigate");
         YQL_ENSURE(ev->Get()->Request);
         if (ev->Get()->Request->Cookie < QueryId) {
+            LOG_I("unexpected return #2");
             return;
         }
 
@@ -2936,21 +2946,33 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
         }
 
         if (HasTopicWriteOperations() && !HasTopicWriteId()) {
+            LOG_I("begin request for WriteId");
             Send(MakeTxProxyID(), new TEvTxUserProxy::TEvAllocateTxId, 0, QueryState->QueryId);
         } else {
+            LOG_I("end request for TopicOperations");
             ReplySuccess();
+
+            LOG_I("after ReplySuccess");
         }
     }
 
     void Handle(TEvTxUserProxy::TEvAllocateTxIdResult::TPtr& ev) {
+        LOG_I("end request for WriteId");
         if (CurrentStateFunc() != &TThis::ExecuteState || ev->Cookie < QueryId) {
+            LOG_I("unexpected return #1");
             return;
         }
 
         YQL_ENSURE(QueryState);
         YQL_ENSURE(QueryState->GetAction() == NKikimrKqp::QUERY_ACTION_TOPIC);
+
         SetTopicWriteId(NLongTxService::TLockHandle(ev->Get()->TxId, TActivationContext::ActorSystem()));
+
+        LOG_I("end request for TopicOperations");
         ReplySuccess();
+
+        LOG_I("current state: " << CurrentStateFuncName());
+        LOG_I("after ReplySuccess");
     }
 
     bool HasTopicWriteOperations() const {

ydb/core/kqp/topics/kqp_topics.cpp

@@ -412,6 +412,8 @@ bool TTopicOperations::ProcessSchemeCacheNavigate(const NSchemeCache::TSchemeCac
                     p->second.SetTabletId(partition.GetTabletId());
                 }
             }
+
+            CachedNavigateResult_[path] = result;
         } else {
             builder << "Topic '" << JoinPath(result.Path) << "' is missing";
 
@@ -428,9 +430,34 @@ bool TTopicOperations::ProcessSchemeCacheNavigate(const NSchemeCache::TSchemeCac
     return true;
 }
 
-bool TTopicOperations::HasThisPartitionAlreadyBeenAdded(const TString& topic, ui32 partitionId) const
+bool TTopicOperations::HasThisPartitionAlreadyBeenAdded(const TString& topicPath, ui32 partitionId)
 {
-    return Operations_.contains({topic, partitionId});
+    if (Operations_.contains({topicPath, partitionId})) {
+        return true;
+    }
+    if (!CachedNavigateResult_.contains(topicPath)) {
+        return false;
+    }
+
+    const NSchemeCache::TSchemeCacheNavigate::TEntry& entry =
+        CachedNavigateResult_.at(topicPath);
+    const NKikimrSchemeOp::TPersQueueGroupDescription& description =
+        entry.PQGroupInfo->Description;
+
+    TString path = CanonizePath(entry.Path);
+    Y_ABORT_UNLESS(path == topicPath,
+                   "path=%s, topicPath=%s",
+                   path.data(), topicPath.data());
+
+    for (const auto& partition : description.GetPartitions()) {
+        if (partition.GetPartitionId() == partitionId) {
+            TTopicPartition key{topicPath, partitionId};
+            Operations_[key].SetTabletId(partition.GetTabletId());
+            return true;
+        }
+    }
+
+    return false;
 }
 
 void TTopicOperations::BuildTopicTxs(TTopicOperationTransactions& txs)

ydb/core/kqp/topics/kqp_topics.h

@@ -159,7 +159,7 @@ class TTopicOperations {
 
     size_t GetSize() const;
 
-    bool HasThisPartitionAlreadyBeenAdded(const TString& topic, ui32 partitionId) const;
+    bool HasThisPartitionAlreadyBeenAdded(const TString& topic, ui32 partitionId);
 
 private:
     THashMap<TTopicPartition, TTopicPartitionOperations, TTopicPartition::THash> Operations_;
@@ -168,6 +168,8 @@ class TTopicOperations {
 
     TMaybe<TString> Consumer_;
     NLongTxService::TLockHandle WriteId_;
+
+    THashMap<TString, NSchemeCache::TSchemeCacheNavigate::TEntry> CachedNavigateResult_;
 };
 
 }