improve lock scope

apache · Mar 15, 2024 · ca36be2 · ca36be2
1 parent 8a9fd9b
commit ca36be2
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 24 deletions.
diff --git a/cpp-ch/local-engine/Parser/MergeTreeRelParser.cpp b/cpp-ch/local-engine/Parser/MergeTreeRelParser.cpp
@@ -137,10 +137,8 @@ MergeTreeRelParser::parseReadRel(
                 buildMergeTreeSettings(merge_tree_table.table_configs));
             return custom_storage_merge_tree;
         });
-    {
-        auto lock = storage->lockForAlter(context->getSettingsRef().lock_acquire_timeout);
-        restoreMetaData(storage->getStoragePolicy()->getAnyDisk(), merge_tree_table);
-    }
+
+    restoreMetaData(storage, merge_tree_table, context);
     for (const auto & [name, sizes] : storage->getColumnSizes())
         column_sizes[name] = sizes.data_compressed;
     query_context.storage_snapshot = std::make_shared<StorageSnapshot>(*storage, metadata);

diff --git a/cpp-ch/local-engine/Storages/Mergetree/MetaDataHelper.cpp b/cpp-ch/local-engine/Storages/Mergetree/MetaDataHelper.cpp
@@ -17,6 +17,8 @@
 #include "MetaDataHelper.h"
 #include <filesystem>
 
+#include <Disks/ObjectStorages/MetadataStorageFromDisk.h>
+
 using namespace DB;
 
 namespace local_engine
@@ -41,33 +43,51 @@ std::unordered_map<String, String> extractPartMetaData(ReadBuffer & in)
     return result;
 }
 
-void restoreMetaData(DiskPtr data_disk, const MergeTreeTable & mergeTreeTable)
+void restoreMetaData(CustomStorageMergeTreePtr & storage, const MergeTreeTable & mergeTreeTable, ContextPtr & context)
 {
+    auto data_disk = storage->getStoragePolicy()->getAnyDisk();
     if (!data_disk->isRemote())
         return;
-    auto s3 = data_disk->getObjectStorage();
-    MetadataStorageFromDisk * metadata_storage = static_cast<MetadataStorageFromDisk *>(data_disk->getMetadataStorage().get());
+
+    std::unordered_set<String> not_exists_part;
+    DB::MetadataStorageFromDisk * metadata_storage = static_cast<MetadataStorageFromDisk *>(data_disk->getMetadataStorage().get());
     auto metadata_disk = metadata_storage->getDisk();
     auto table_path = std::filesystem::path(mergeTreeTable.relative_path);
-    if (!metadata_disk->exists(table_path))
-        metadata_disk->createDirectories(table_path.generic_string());
     for (const auto & part : mergeTreeTable.getPartNames())
     {
         auto part_path = table_path / part;
-        auto metadata_file_path = part_path / "metadata.gluten";
+        if (!metadata_disk->exists(part_path))
+            not_exists_part.emplace(part);
+    }
+
+    if (not_exists_part.empty())
+        return;
 
-        if (metadata_disk->exists(part_path))
-            continue;
-        else
-            metadata_disk->createDirectories(part_path);
-        auto key = s3->generateObjectKeyForPath(metadata_file_path.generic_string());
-        StoredObject metadata_object(key.serialize());
-        auto part_metadata = extractPartMetaData(*s3->readObject(metadata_object));
-        for (const auto & item : part_metadata)
+    if (auto lock = storage->lockForAlter(context->getSettingsRef().lock_acquire_timeout))
+    {
+        auto s3 = data_disk->getObjectStorage();
+
+        if (!metadata_disk->exists(table_path))
+            metadata_disk->createDirectories(table_path.generic_string());
+
+        for (const auto & part : not_exists_part)
         {
-            auto item_path = part_path / item.first;
-            auto out = metadata_disk->writeFile(item_path);
-            out->write(item.second.data(), item.second.size());
+            auto part_path = table_path / part;
+            auto metadata_file_path = part_path / "metadata.gluten";
+
+            if (metadata_disk->exists(part_path))
+                continue;
+            else
+                metadata_disk->createDirectories(part_path);
+            auto key = s3->generateObjectKeyForPath(metadata_file_path.generic_string());
+            StoredObject metadata_object(key.serialize());
+            auto part_metadata = extractPartMetaData(*s3->readObject(metadata_object));
+            for (const auto & item : part_metadata)
+            {
+                auto item_path = part_path / item.first;
+                auto out = metadata_disk->writeFile(item_path);
+                out->write(item.second.data(), item.second.size());
+            }
         }
     }
 }

diff --git a/cpp-ch/local-engine/Storages/Mergetree/MetaDataHelper.h b/cpp-ch/local-engine/Storages/Mergetree/MetaDataHelper.h
@@ -17,14 +17,13 @@
 
 #pragma once
 
-#include <Disks/ObjectStorages/MetadataStorageFromDisk.h>
 #include <Common/MergeTreeTool.h>
-
+#include <Storages/StorageMergeTreeFactory.h>
 
 namespace local_engine
 {
 
-void restoreMetaData(DB::DiskPtr data_disk, const MergeTreeTable & mergeTreeTable);
+void restoreMetaData(CustomStorageMergeTreePtr & storage, const MergeTreeTable & mergeTreeTable, ContextPtr & context);
 
 }