From 22487de34a920cb5199f8289dfc6951e608a0780 Mon Sep 17 00:00:00 2001
From: kssenii <sumarokovakseniia@mail.ru>
Date: Tue, 2 May 2023 16:09:33 +0200
Subject: [PATCH] Fix

---
 src/Interpreters/Cache/FileSegment.cpp | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/Interpreters/Cache/FileSegment.cpp b/src/Interpreters/Cache/FileSegment.cpp
index b05847a1ce7..300853c4423 100644
--- a/src/Interpreters/Cache/FileSegment.cpp
+++ b/src/Interpreters/Cache/FileSegment.cpp
@@ -56,6 +56,7 @@ FileSegment::FileSegment(
         /// someone will _potentially_ want to download it (after calling getOrSetDownloader()).
         case (State::EMPTY):
         {
+            chassert(key_metadata.lock());
             break;
         }
         /// DOWNLOADED is used either on initial cache metadata load into memory on server startup
@@ -65,6 +66,7 @@ FileSegment::FileSegment(
             reserved_size = downloaded_size = size_;
             chassert(fs::file_size(getPathInLocalCache()) == size_);
             chassert(queue_iterator);
+            chassert(key_metadata.lock());
             break;
         }
         case (State::DETACHED):
@@ -91,8 +93,16 @@ String FileSegment::getPathInLocalCache() const
     return getKeyMetadata()->getFileSegmentPath(*this);
 }
 
-void FileSegment::setDownloadState(State state, const FileSegmentGuard::Lock &)
+void FileSegment::setDownloadState(State state, const FileSegmentGuard::Lock & lock)
 {
+    if (isCompleted(false))
+    {
+        throw Exception(
+            ErrorCodes::LOGICAL_ERROR,
+            "Updating state of file segment is not allowed, because it is already completed ({})",
+            getInfoForLogUnlocked(lock));
+    }
+
     LOG_TEST(log, "Updated state from {} to {}", stateToString(download_state), stateToString(state));
     download_state = state;
 }
@@ -182,12 +192,13 @@ String FileSegment::getOrSetDownloader()
     if (current_downloader.empty())
     {
         const auto caller_id = getCallerId();
-        bool allow_new_downloader = download_state == State::EMPTY || download_state == State::PARTIALLY_DOWNLOADED || !caller_id.starts_with("None");
+        bool allow_new_downloader = download_state == State::EMPTY || download_state == State::PARTIALLY_DOWNLOADED;
         if (!allow_new_downloader)
             return "notAllowed:" + stateToString(download_state);
 
         current_downloader = downloader_id = caller_id;
         setDownloadState(State::DOWNLOADING, lock);
+        chassert(key_metadata.lock());
     }
 
     return current_downloader;