Allow to shrink file segment only to aligned size

src/Interpreters/Cache/FileCache.cpp

@@ -9,6 +9,7 @@
 #include <Interpreters/Cache/LRUFileCachePriority.h>
 #include <Interpreters/Cache/SLRUFileCachePriority.h>
 #include <Interpreters/Cache/EvictionCandidates.h>
+#include <Interpreters/Cache/FileCacheUtils.h>
 #include <Interpreters/Context.h>
 #include <base/hex.h>
 #include <Common/callOnce.h>
@@ -53,16 +54,6 @@ namespace ErrorCodes
 
 namespace
 {
-    size_t roundDownToMultiple(size_t num, size_t multiple)
-    {
-        return (num / multiple) * multiple;
-    }
-
-    size_t roundUpToMultiple(size_t num, size_t multiple)
-    {
-        return roundDownToMultiple(num + multiple - 1, multiple);
-    }
-
     std::string getCommonUserID()
     {
         auto user_from_context = DB::Context::getGlobalContextInstance()->getFilesystemCacheUser();
@@ -601,8 +592,8 @@ FileCache::getOrSet(
     /// 2. max_file_segments_limit
     FileSegment::Range result_range = initial_range;
 
-    const auto aligned_offset = roundDownToMultiple(initial_range.left, boundary_alignment);
-    auto aligned_end_offset = std::min(roundUpToMultiple(initial_range.right + 1, boundary_alignment), file_size) - 1;
+    const auto aligned_offset = FileCacheUtils::roundDownToMultiple(initial_range.left, boundary_alignment);
+    auto aligned_end_offset = std::min(FileCacheUtils::roundUpToMultiple(initial_range.right + 1, boundary_alignment), file_size) - 1;
 
     chassert(aligned_offset <= initial_range.left);
     chassert(aligned_end_offset >= initial_range.right);

src/Interpreters/Cache/FileCache.h

@@ -161,6 +161,8 @@ public:
 
     size_t getMaxFileSegmentSize() const { return max_file_segment_size; }
 
+    size_t getBoundaryAlignment() const { return boundary_alignment; }
+
     bool tryReserve(
         FileSegment & file_segment,
         size_t size,

src/Interpreters/Cache/FileCacheUtils.h

@@ -0,0 +1,17 @@
+#pragma once
+#include <Core/Types.h>
+
+namespace FileCacheUtils
+{
+
+static size_t roundDownToMultiple(size_t num, size_t multiple)
+{
+    return (num / multiple) * multiple;
+}
+
+static size_t roundUpToMultiple(size_t num, size_t multiple)
+{
+    return roundDownToMultiple(num + multiple - 1, multiple);
+}
+
+}

src/Interpreters/Cache/FileSegment.cpp

@@ -4,6 +4,7 @@
 #include <IO/Operators.h>
 #include <IO/WriteBufferFromString.h>
 #include <Interpreters/Cache/FileCache.h>
+#include <Interpreters/Cache/FileCacheUtils.h>
 #include <base/getThreadId.h>
 #include <base/hex.h>
 #include <Common/CurrentThread.h>
@@ -85,7 +86,6 @@ FileSegment::FileSegment(
             break;
         }
         /// DOWNLOADED is used either on initial cache metadata load into memory on server startup
-        /// or on shrinkFileSegmentToDownloadedSize() -- when file segment object is updated.
         case (State::DOWNLOADED):
         {
             reserved_size = downloaded_size = size_;
@@ -629,6 +629,46 @@ void FileSegment::completePartAndResetDownloader()
     LOG_TEST(log, "Complete batch. ({})", getInfoForLogUnlocked(lk));
 }
 
+void FileSegment::shrinkFileSegmentToDownloadedSize(const LockedKey & locked_key, const FileSegmentGuard::Lock & lock)
+{
+    if (downloaded_size == range().size())
+    {
+        /// Nothing to resize;
+        return;
+    }
+
+    if (!locked_key.isLastOwnerOfFileSegment(offset()))
+    {
+        throw Exception(
+            ErrorCodes::LOGICAL_ERROR,
+            "Shrinking of file segment can  be done only by the last holder: {}",
+            getInfoForLog());
+    }
+
+    const size_t aligned_downloaded_size = FileCacheUtils::roundUpToMultiple(downloaded_size, cache->getBoundaryAlignment());
+
+    chassert(aligned_downloaded_size <= range().size());
+    chassert(aligned_downloaded_size >= downloaded_size);
+
+    if (aligned_downloaded_size == range().size())
+    {
+        /// Nothing to resize;
+        return;
+    }
+
+    if (downloaded_size == aligned_downloaded_size)
+        setDownloadState(State::DOWNLOADED, lock);
+    else
+        setDownloadState(State::PARTIALLY_DOWNLOADED, lock);
+
+    segment_range.right = segment_range.left + aligned_downloaded_size - 1;
+
+    const size_t diff = reserved_size - downloaded_size;
+    chassert(reserved_size >= downloaded_size);
+    if (diff)
+        getQueueIterator()->decrementSize(diff);
+}
+
 void FileSegment::complete(bool allow_background_download)
 {
     ProfileEventTimeIncrement<Microseconds> watch(ProfileEvents::FileSegmentCompleteMicroseconds);
@@ -716,8 +756,7 @@ void FileSegment::complete(bool allow_background_download)
 
                 if (!added_to_download_queue)
                 {
-                    locked_key->shrinkFileSegmentToDownloadedSize(offset(), segment_lock);
-                    setDetachedState(segment_lock); /// See comment below.
+                    shrinkFileSegmentToDownloadedSize(*locked_key, segment_lock);
                 }
             }
             break;
@@ -747,11 +786,7 @@ void FileSegment::complete(bool allow_background_download)
                     /// but current file segment should remain PARRTIALLY_DOWNLOADED_NO_CONTINUATION and with detached state,
                     /// because otherwise an invariant that getOrSet() returns a contiguous range of file segments will be broken
                     /// (this will be crucial for other file segment holder, not for current one).
-                    locked_key->shrinkFileSegmentToDownloadedSize(offset(), segment_lock);
-
-                    /// We mark current file segment with state DETACHED, even though the data is still in cache
-                    /// (but a separate file segment) because is_last_holder is satisfied, so it does not matter.
-                    setDetachedState(segment_lock);
+                    shrinkFileSegmentToDownloadedSize(*locked_key, segment_lock);
                 }
             }
             break;

src/Interpreters/Cache/FileSegment.h

@@ -239,6 +239,7 @@ private:
 
     void setDownloadedUnlocked(const FileSegmentGuard::Lock &);
     void setDownloadFailedUnlocked(const FileSegmentGuard::Lock &);
+    void shrinkFileSegmentToDownloadedSize(const LockedKey &, const FileSegmentGuard::Lock &);
 
     void assertNotDetached() const;
     void assertNotDetachedUnlocked(const FileSegmentGuard::Lock &) const;

src/Interpreters/Cache/Metadata.cpp

@@ -990,42 +990,6 @@ KeyMetadata::iterator LockedKey::removeFileSegmentImpl(
     return key_metadata->erase(it);
 }
 
-void LockedKey::shrinkFileSegmentToDownloadedSize(
-    size_t offset,
-    const FileSegmentGuard::Lock & segment_lock)
-{
-    /**
-     * In case file was partially downloaded and it's download cannot be continued
-     * because of no space left in cache, we need to be able to cut file segment's size to downloaded_size.
-     */
-
-    auto file_segment_metadata = getByOffset(offset);
-    const auto & file_segment = file_segment_metadata->file_segment;
-    chassert(file_segment->assertCorrectnessUnlocked(segment_lock));
-
-    const size_t downloaded_size = file_segment->getDownloadedSize();
-    if (downloaded_size == file_segment->range().size())
-    {
-        throw Exception(
-            ErrorCodes::LOGICAL_ERROR,
-            "Nothing to reduce, file segment fully downloaded: {}",
-            file_segment->getInfoForLogUnlocked(segment_lock));
-    }
-
-    chassert(file_segment->reserved_size >= downloaded_size);
-    int64_t diff = file_segment->reserved_size - downloaded_size;
-
-    file_segment_metadata->file_segment = std::make_shared<FileSegment>(
-        getKey(), offset, downloaded_size, FileSegment::State::DOWNLOADED,
-        CreateFileSegmentSettings(file_segment->getKind()), false,
-        file_segment->cache, key_metadata, file_segment->queue_iterator);
-
-    if (diff)
-        file_segment_metadata->getQueueIterator()->decrementSize(diff);
-
-    chassert(file_segment_metadata->file_segment->assertCorrectnessUnlocked(segment_lock));
-}
-
 bool LockedKey::addToDownloadQueue(size_t offset, const FileSegmentGuard::Lock &)
 {
     auto it = key_metadata->find(offset);

src/Interpreters/Cache/Metadata.h

@@ -325,8 +325,6 @@ struct LockedKey : private boost::noncopyable
         bool can_be_broken = false,
         bool invalidate_queue_entry = true);
 
-    void shrinkFileSegmentToDownloadedSize(size_t offset, const FileSegmentGuard::Lock &);
-
     bool addToDownloadQueue(size_t offset, const FileSegmentGuard::Lock &);
 
     bool isLastOwnerOfFileSegment(size_t offset) const;