Tune replay of PathStat (#10311)

ydb-platform · Oct 11, 2024 · 38747d7 · 38747d7
1 parent 65aa640
commit 38747d7
Showing 1 changed file with 40 additions and 1 deletion.
diff --git a/ydb/library/yql/providers/yt/gateway/qplayer/yql_yt_qplayer_gateway.cpp b/ydb/library/yql/providers/yt/gateway/qplayer/yql_yt_qplayer_gateway.cpp
@@ -25,6 +25,7 @@ const TString YtGateway_GetFolder = "YtGateway_GetFolder";
 const TString YtGateway_GetFolders = "YtGateway_GetFolders";
 const TString YtGateway_ResolveLinks = "YtGateway_ResolveLinks";
 const TString YtGateway_PathStat = "YtGateway_PathStat";
+const TString YtGateway_PathStatMissing = "YtGateway_PathStatMissing";
 
 TString MakeHash(const TString& str) {
     SHA256_CTX sha;
@@ -777,6 +778,7 @@ class TGateway : public IYtGateway {
                     throw yexception() << "Missing replay data";
                 }
 
+                PathStatKeys_.emplace(key);
                 auto valueNode = NYT::NodeFromYsonString(TStringBuf(item->Value));
                 DeserializePathStat(valueNode, res, index);
             }
@@ -813,6 +815,21 @@ class TGateway : public IYtGateway {
 
             for (ui32 index = 0; index < options.Paths().size(); ++index) {
                 const auto& key = MakePathStatKey(options.Cluster(), options.Extended(), options.Paths()[index]);
+                bool allow = false;
+                if (PathStatKeys_.contains(key)) {
+                    allow = true;
+                } else {
+                    auto missingItem = QContext_.GetReader()->Get({YtGateway_PathStatMissing, key}).GetValueSync();
+                    if (!missingItem) {
+                        allow = true;
+                        PathStatKeys_.emplace(key);
+                    }
+                }
+
+                if (!allow) {
+                    return res;
+                }
+
                 auto item = QContext_.GetReader()->Get({YtGateway_PathStat, key}).GetValueSync();
                 if (!item) {
                     return res;
@@ -826,7 +843,28 @@ class TGateway : public IYtGateway {
             return res;
         }
 
-        return Inner_->TryPathStat(std::move(options));
+        auto optionsDup = options;
+        auto res = Inner_->TryPathStat(std::move(options));
+        if (!QContext_.CanWrite()) {
+            return res;
+        }
+
+        if (!res.Success()) {
+            for (ui32 index = 0; index < optionsDup.Paths().size(); ++index) {
+                const auto& key = MakePathStatKey(optionsDup.Cluster(), optionsDup.Extended(), optionsDup.Paths()[index]);
+                QContext_.GetWriter()->Put({YtGateway_PathStatMissing, key}, "1").GetValueSync();
+            }
+
+            return res;
+        }
+
+        for (ui32 index = 0; index < optionsDup.Paths().size(); ++index) {
+            const auto& key = MakePathStatKey(optionsDup.Cluster(), optionsDup.Extended(), optionsDup.Paths()[index]);
+            auto value = SerializePathStat(res, index);
+            QContext_.GetWriter()->Put({YtGateway_PathStat, key}, value).GetValueSync();
+        }
+
+        return res;
     }
 
     bool TryParseYtUrl(const TString& url, TString* cluster, TString* path) const final {
@@ -882,6 +920,7 @@ class TGateway : public IYtGateway {
     const TQContext QContext_;
     const TIntrusivePtr<IRandomProvider> RandomProvider_;
     const TFileStoragePtr FileStorage_;
+    THashSet<TString> PathStatKeys_;
 };
 
 }