Merge pull request #17737 from ClickHouse/fix_segfault_in_distributed_out_stream

Fix segfault when 'not enough space'

(cherry picked from commit 0eec52b1)

src/Disks/StoragePolicy.cpp

@@ -30,6 +30,7 @@ namespace ErrorCodes
     extern const int UNKNOWN_POLICY;
     extern const int UNKNOWN_VOLUME;
     extern const int LOGICAL_ERROR;
+    extern const int NOT_ENOUGH_SPACE;
 }
 
 
@@ -210,6 +211,14 @@ ReservationPtr StoragePolicy::reserve(UInt64 bytes) const
 }
 
 
+ReservationPtr StoragePolicy::reserveAndCheck(UInt64 bytes) const
+{
+    if (auto res = reserve(bytes, 0))
+        return res;
+    throw Exception(ErrorCodes::NOT_ENOUGH_SPACE, "Cannot reserve {}, not enough space", ReadableSize(bytes));
+}
+
+
 ReservationPtr StoragePolicy::makeEmptyReservationOnLargestDisk() const
 {
     UInt64 max_space = 0;
@@ -226,7 +235,14 @@ ReservationPtr StoragePolicy::makeEmptyReservationOnLargestDisk() const
             }
         }
     }
-    return max_disk->reserve(0);
+    auto reservation = max_disk->reserve(0);
+    if (!reservation)
+    {
+        /// I'm not sure if it's really a logical error, but exception message
+        /// "Cannot reserve 0 bytes" looks too strange to throw it with another exception code.
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "Cannot reserve 0 bytes");
+    }
+    return reservation;
 }
 
 

src/Disks/StoragePolicy.h

@@ -61,10 +61,13 @@ public:
 
     const String & getName() const { return name; }
 
-    /// Returns valid reservation or null
+    /// Returns valid reservation or nullptr
     ReservationPtr reserve(UInt64 bytes) const;
 
-    /// Reserve space on any volume with index > min_volume_index
+    /// Reserves space on any volume or throws
+    ReservationPtr reserveAndCheck(UInt64 bytes) const;
+
+    /// Reserves space on any volume with index > min_volume_index or returns nullptr
     ReservationPtr reserve(UInt64 bytes, size_t min_volume_index) const;
 
     /// Find volume index, which contains disk

src/Storages/Distributed/DistributedBlockOutputStream.cpp

@@ -567,7 +567,7 @@ void DistributedBlockOutputStream::writeToShard(const Block & block, const std::
     /// and keep monitor thread out from reading incomplete data
     std::string first_file_tmp_path{};
 
-    auto reservation = storage.getStoragePolicy()->reserve(block.bytes());
+    auto reservation = storage.getStoragePolicy()->reserveAndCheck(block.bytes());
     auto disk = reservation->getDisk()->getPath();
     auto data_path = storage.getRelativeDataPath();
 

src/Storages/MergeTree/MergeTreeData.cpp

@@ -3140,8 +3140,7 @@ inline ReservationPtr checkAndReturnReservation(UInt64 expected_size, Reservatio
 ReservationPtr MergeTreeData::reserveSpace(UInt64 expected_size) const
 {
     expected_size = std::max(RESERVATION_MIN_ESTIMATION_SIZE, expected_size);
-    auto reservation = getStoragePolicy()->reserve(expected_size);
-    return checkAndReturnReservation(expected_size, std::move(reservation));
+    return getStoragePolicy()->reserveAndCheck(expected_size);
 }
 
 ReservationPtr MergeTreeData::reserveSpace(UInt64 expected_size, SpacePtr space)