Merge pull request #55322 from ClickHouse/cherrypick/23.8/20868f3b656…

…466021f57ccabbe15c1c278ebc826

Cherry pick #55252 to 23.8: Fix for background download in fs cache

src/Interpreters/Cache/FileCache.cpp

@@ -519,7 +519,7 @@ KeyMetadata::iterator FileCache::addFileSegment(
         result_state = state;
     }
 
-    auto file_segment = std::make_shared<FileSegment>(key, offset, size, result_state, settings, this, locked_key.getKeyMetadata());
+    auto file_segment = std::make_shared<FileSegment>(key, offset, size, result_state, settings, background_download_threads > 0, this, locked_key.getKeyMetadata());
     auto file_segment_metadata = std::make_shared<FileSegmentMetadata>(std::move(file_segment));
 
     auto [file_segment_metadata_it, inserted] = locked_key.getKeyMetadata()->emplace(offset, file_segment_metadata);
@@ -1016,6 +1016,7 @@ void FileCache::loadMetadataForKeys(const fs::path & keys_dir)
                     auto file_segment = std::make_shared<FileSegment>(key, offset, size,
                                                                       FileSegment::State::DOWNLOADED,
                                                                       CreateFileSegmentSettings(segment_kind),
+                                                                      false,
                                                                       this,
                                                                       key_metadata,
                                                                       cache_it);

src/Interpreters/Cache/FileCache.h

@@ -158,7 +158,7 @@ class FileCache : private boost::noncopyable
     const size_t max_file_segment_size;
     const size_t bypass_cache_threshold = 0;
     const size_t boundary_alignment;
-    const size_t background_download_threads;
+    const size_t background_download_threads; /// 0 means background download is disabled.
     const size_t metadata_download_threads;
 
     Poco::Logger * log;

src/Interpreters/Cache/FileSegment.cpp

@@ -44,13 +44,15 @@ FileSegment::FileSegment(
         size_t size_,
         State download_state_,
         const CreateFileSegmentSettings & settings,
+        bool background_download_enabled_,
         FileCache * cache_,
         std::weak_ptr<KeyMetadata> key_metadata_,
         Priority::Iterator queue_iterator_)
     : file_key(key_)
     , segment_range(offset_, offset_ + size_ - 1)
     , segment_kind(settings.kind)
     , is_unbound(settings.unbounded)
+    , background_download_enabled(background_download_enabled_)
     , download_state(download_state_)
     , key_metadata(key_metadata_)
     , queue_iterator(queue_iterator_)
@@ -539,6 +541,12 @@ void FileSegment::setDownloadedUnlocked(const FileSegmentGuard::Lock &)
     chassert(fs::file_size(getPathInLocalCache()) == downloaded_size);
 }
 
+void FileSegment::setDownloadFailed()
+{
+    auto lock = lockFileSegment();
+    setDownloadFailedUnlocked(lock);
+}
+
 void FileSegment::setDownloadFailedUnlocked(const FileSegmentGuard::Lock & lock)
 {
     LOG_INFO(log, "Setting download as failed: {}", getInfoForLogUnlocked(lock));
@@ -652,7 +660,7 @@ void FileSegment::complete()
 
             if (is_last_holder)
             {
-                if (remote_file_reader)
+                if (background_download_enabled && remote_file_reader)
                 {
                     LOG_TEST(
                         log, "Submitting file segment for background download "

src/Interpreters/Cache/FileSegment.h

@@ -110,6 +110,7 @@ friend class FileCache; /// Because of reserved_size in tryReserve().
         size_t size_,
         State download_state_,
         const CreateFileSegmentSettings & create_settings = {},
+        bool background_download_enabled_ = false,
         FileCache * cache_ = nullptr,
         std::weak_ptr<KeyMetadata> key_metadata_ = std::weak_ptr<KeyMetadata>(),
         Priority::Iterator queue_iterator_ = Priority::Iterator{});
@@ -262,6 +263,8 @@ friend class FileCache; /// Because of reserved_size in tryReserve().
 
     void setDownloadedSize(size_t delta);
 
+    void setDownloadFailed();
+
 private:
     String getDownloaderUnlocked(const FileSegmentGuard::Lock &) const;
     bool isDownloaderUnlocked(const FileSegmentGuard::Lock & segment_lock) const;
@@ -290,6 +293,7 @@ friend class FileCache; /// Because of reserved_size in tryReserve().
     /// Size of the segment is not known until it is downloaded and
     /// can be bigger than max_file_segment_size.
     const bool is_unbound = false;
+    const bool background_download_enabled;
 
     std::atomic<State> download_state;
     DownloaderId downloader_id; /// The one who prepares the download

src/Interpreters/Cache/Metadata.cpp

@@ -520,65 +520,74 @@ void CacheMetadata::downloadThreadFunc()
 
         CurrentMetrics::sub(CurrentMetrics::FilesystemCacheDownloadQueueElements);
 
-        FileSegmentsHolderPtr holder;
         try
         {
+            FileSegmentsHolderPtr holder;
+            try
             {
-                auto locked_key = lockKeyMetadata(key, KeyNotFoundPolicy::RETURN_NULL);
-                if (!locked_key)
-                    continue;
+                {
+                    auto locked_key = lockKeyMetadata(key, KeyNotFoundPolicy::RETURN_NULL);
+                    if (!locked_key)
+                        continue;
 
-                auto file_segment_metadata = locked_key->tryGetByOffset(offset);
-                if (!file_segment_metadata || file_segment_metadata->evicting())
-                    continue;
+                    auto file_segment_metadata = locked_key->tryGetByOffset(offset);
+                    if (!file_segment_metadata || file_segment_metadata->evicting())
+                        continue;
+
+                    auto file_segment = file_segment_weak.lock();
+
+                    if (!file_segment
+                        || file_segment != file_segment_metadata->file_segment
+                        || file_segment->state() != FileSegment::State::PARTIALLY_DOWNLOADED)
+                        continue;
+
+                    holder = std::make_unique<FileSegmentsHolder>(FileSegments{file_segment});
+                }
 
-                auto file_segment = file_segment_weak.lock();
+                auto & file_segment = holder->front();
 
-                if (!file_segment
-                    || file_segment != file_segment_metadata->file_segment
-                    || file_segment->state() != FileSegment::State::PARTIALLY_DOWNLOADED)
+                if (file_segment.getOrSetDownloader() != FileSegment::getCallerId())
                     continue;
 
-                holder = std::make_unique<FileSegmentsHolder>(FileSegments{file_segment});
-            }
+                chassert(file_segment.getDownloadedSize() != file_segment.range().size());
+                chassert(file_segment.assertCorrectness());
 
-            downloadImpl(holder->front(), memory);
-        }
-        catch (...)
-        {
-            if (holder)
-            {
-                const auto & file_segment = holder->front();
-                LOG_ERROR(
-                    log, "Error during background download of {}:{} ({}): {}",
-                    file_segment.key(), file_segment.offset(),
-                    file_segment.getInfoForLog(), getCurrentExceptionMessage(true));
+                downloadImpl(file_segment, memory);
             }
-            else
+            catch (...)
             {
-                tryLogCurrentException(__PRETTY_FUNCTION__);
-                chassert(false);
+                if (holder)
+                {
+                    auto & file_segment = holder->front();
+                    file_segment.setDownloadFailed();
+
+                    LOG_ERROR(
+                        log, "Error during background download of {}:{} ({}): {}",
+                        file_segment.key(), file_segment.offset(),
+                        file_segment.getInfoForLog(), getCurrentExceptionMessage(true));
+                }
+                else
+                {
+                    tryLogCurrentException(__PRETTY_FUNCTION__);
+                    chassert(false);
+                }
             }
         }
+        catch (...)
+        {
+            tryLogCurrentException(__PRETTY_FUNCTION__);
+            chassert(false);
+        }
     }
 }
 
 void CacheMetadata::downloadImpl(FileSegment & file_segment, std::optional<Memory<>> & memory)
 {
-    chassert(file_segment.assertCorrectness());
-
-    if (file_segment.getOrSetDownloader() != FileSegment::getCallerId())
-        return;
-
-    if (file_segment.getDownloadedSize() == file_segment.range().size())
-        throw Exception(ErrorCodes::LOGICAL_ERROR, "File segment is already fully downloaded");
-
     LOG_TEST(
         log, "Downloading {} bytes for file segment {}",
         file_segment.range().size() - file_segment.getDownloadedSize(), file_segment.getInfoForLog());
 
     auto reader = file_segment.getRemoteFileReader();
-
     if (!reader)
     {
         throw Exception(
@@ -808,7 +817,7 @@ void LockedKey::shrinkFileSegmentToDownloadedSize(
 
     metadata->file_segment = std::make_shared<FileSegment>(
         getKey(), offset, downloaded_size, FileSegment::State::DOWNLOADED,
-        CreateFileSegmentSettings(file_segment->getKind()),
+        CreateFileSegmentSettings(file_segment->getKind()), false,
         file_segment->cache, key_metadata, file_segment->queue_iterator);
 
     if (diff)