Merge pull request #40608 from ClickHouse/kssenii-patch-4

Fix stress test after #40420

src/Common/FileSegment.cpp

@@ -662,6 +662,8 @@ void FileSegment::assertDetachedStatus(std::lock_guard<std::mutex> & segment_loc
 
 FileSegmentPtr FileSegment::getSnapshot(const FileSegmentPtr & file_segment, std::lock_guard<std::mutex> & /* cache_lock */)
 {
+    std::lock_guard segment_lock(file_segment->mutex);
+
     auto snapshot = std::make_shared<FileSegment>(
         file_segment->offset(),
         file_segment->range().size(),
@@ -671,8 +673,8 @@ FileSegmentPtr FileSegment::getSnapshot(const FileSegmentPtr & file_segment, std
 
     snapshot->hits_count = file_segment->getHitsCount();
     snapshot->ref_count = file_segment.use_count();
-    snapshot->downloaded_size = file_segment->getDownloadedSize();
-    snapshot->download_state = file_segment->state();
+    snapshot->downloaded_size = file_segment->getDownloadedSize(segment_lock);
+    snapshot->download_state = file_segment->download_state;
     snapshot->is_persistent = file_segment->isPersistent();
 
     return snapshot;
@@ -818,15 +820,25 @@ void FileSegmentRangeWriter::completeFileSegment(FileSegment & file_segment)
 
     if (file_segment.getDownloadedSize() > 0)
     {
-        /// file_segment->complete(DOWNLOADED) is not enough, because file segment capacity
-        /// was initially set with a margin as `max_file_segment_size`. => We need to always
-        /// resize to actual size after download finished.
         file_segment.getOrSetDownloader();
 
-        assert(file_segment.downloaded_size <= file_segment.range().size());
-        file_segment.segment_range = FileSegment::Range(
-            file_segment.segment_range.left, file_segment.segment_range.left + file_segment.downloaded_size - 1);
-        file_segment.reserved_size = file_segment.downloaded_size;
+        {
+            /// file_segment->complete(DOWNLOADED) is not enough, because file segment capacity
+            /// was initially set with a margin as `max_file_segment_size`. => We need to always
+            /// resize to actual size after download finished.
+
+            /// Current file segment is downloaded as a part of write-through cache
+            /// and therefore cannot be concurrently accessed. Nevertheless, it can be
+            /// accessed by cache system tables if someone read from them,
+            /// therefore we need a mutex.
+            std::lock_guard segment_lock(file_segment.mutex);
+
+            assert(file_segment.downloaded_size <= file_segment.range().size());
+            file_segment.segment_range = FileSegment::Range(
+                file_segment.segment_range.left,
+                file_segment.segment_range.left + file_segment.downloaded_size - 1);
+            file_segment.reserved_size = file_segment.downloaded_size;
+        }
 
         file_segment.completeWithState(FileSegment::State::DOWNLOADED);