From b4ef10ad1c075b5c16cafb911f67d4623cd40ceb Mon Sep 17 00:00:00 2001
From: avogar <avogar@clickhouse.com>
Date: Thu, 5 Sep 2024 08:52:22 +0000
Subject: [PATCH] Make better

---
 .../IO/WriteBufferFromAzureBlobStorage.cpp    | 30 +++++--------------
 .../IO/WriteBufferFromAzureBlobStorage.h      |  1 -
 src/IO/WriteBufferFromS3.cpp                  | 21 ++-----------
 src/IO/WriteBufferFromS3.h                    |  1 -
 4 files changed, 10 insertions(+), 43 deletions(-)

diff --git a/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp b/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp
index dcd625f7ee0..cbcfe0bdb97 100644
--- a/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp
+++ b/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp
@@ -59,7 +59,7 @@ WriteBufferFromAzureBlobStorage::WriteBufferFromAzureBlobStorage(
     const WriteSettings & write_settings_,
     std::shared_ptr<const AzureBlobStorage::RequestSettings> settings_,
     ThreadPoolCallbackRunnerUnsafe<void> schedule_)
-    : WriteBufferFromFileBase(buf_size_, nullptr, 0)
+    : WriteBufferFromFileBase(std::min(buf_size_, static_cast<size_t>(DBMS_DEFAULT_BUFFER_SIZE)), nullptr, 0)
     , log(getLogger("WriteBufferFromAzureBlobStorage"))
     , buffer_allocation_policy(createBufferAllocationPolicy(*settings_))
     , max_single_part_upload_size(settings_->max_single_part_upload_size)
@@ -248,30 +248,14 @@ void WriteBufferFromAzureBlobStorage::allocateBuffer()
     buffer_allocation_policy->nextBuffer();
     chassert(0 == hidden_size);
 
-    auto size = buffer_allocation_policy->getBufferSize();
-
-    if (buffer_allocation_policy->getBufferNumber() == 1)
-    {
-        allocateFirstBuffer();
-        return;
-    }
-
-    memory = Memory(size);
-    WriteBuffer::set(memory.data(), memory.size());
-}
-
-void WriteBufferFromAzureBlobStorage::allocateFirstBuffer()
-{
     /// First buffer was already allocated in BufferWithOwnMemory constructor with buffer size provided in constructor.
     /// It will be reallocated in subsequent nextImpl calls up to the desired buffer size from buffer_allocation_policy.
-    /// But it may happen that buffer size provided in constructor is larger then desired buffer size from buffer_allocation_policy.
-    /// Resize memory in this case to the desired size.
-    const auto max_first_buffer = buffer_allocation_policy->getBufferSize();
-    if (memory.size() > max_first_buffer)
-    {
-        memory.resize(max_first_buffer);
-        WriteBuffer::set(memory.data(), memory.size());
-    }
+    if (buffer_allocation_policy->getBufferNumber() == 1)
+        return;
+
+    auto size = buffer_allocation_policy->getBufferSize();
+    memory = Memory(size);
+    WriteBuffer::set(memory.data(), memory.size());
 }
 
 void WriteBufferFromAzureBlobStorage::detachBuffer()
diff --git a/src/Disks/IO/WriteBufferFromAzureBlobStorage.h b/src/Disks/IO/WriteBufferFromAzureBlobStorage.h
index c2d65928cfa..3ee497c4e44 100644
--- a/src/Disks/IO/WriteBufferFromAzureBlobStorage.h
+++ b/src/Disks/IO/WriteBufferFromAzureBlobStorage.h
@@ -53,7 +53,6 @@ private:
     void detachBuffer();
     void reallocateFirstBuffer();
     void allocateBuffer();
-    void allocateFirstBuffer();
     void hidePartialData();
     void setFakeBufferWhenPreFinalized();
 
diff --git a/src/IO/WriteBufferFromS3.cpp b/src/IO/WriteBufferFromS3.cpp
index d88e393e4e1..424708cf2c8 100644
--- a/src/IO/WriteBufferFromS3.cpp
+++ b/src/IO/WriteBufferFromS3.cpp
@@ -96,7 +96,7 @@ WriteBufferFromS3::WriteBufferFromS3(
     std::optional<std::map<String, String>> object_metadata_,
     ThreadPoolCallbackRunnerUnsafe<void> schedule_,
     const WriteSettings & write_settings_)
-    : WriteBufferFromFileBase(buf_size_, nullptr, 0)
+    : WriteBufferFromFileBase(std::min(buf_size_, static_cast<size_t>(DBMS_DEFAULT_BUFFER_SIZE)), nullptr, 0)
     , bucket(bucket_)
     , key(key_)
     , request_settings(request_settings_)
@@ -352,30 +352,15 @@ void WriteBufferFromS3::allocateBuffer()
     buffer_allocation_policy->nextBuffer();
     chassert(0 == hidden_size);
 
+    /// First buffer was already allocated in BufferWithOwnMemory constructor with provided in constructor buffer size.
+    /// It will be reallocated in subsequent nextImpl calls up to the desired buffer size from buffer_allocation_policy.
     if (buffer_allocation_policy->getBufferNumber() == 1)
-    {
-        allocateFirstBuffer();
         return;
-    }
 
     memory = Memory(buffer_allocation_policy->getBufferSize());
     WriteBuffer::set(memory.data(), memory.size());
 }
 
-void WriteBufferFromS3::allocateFirstBuffer()
-{
-    /// First buffer was already allocated in BufferWithOwnMemory constructor with provided in constructor buffer size.
-    /// It will be reallocated in subsequent nextImpl calls up to the desired buffer size from buffer_allocation_policy.
-    /// But it may happen that buffer size provided in constructor is larger then desired buffer size from buffer_allocation_policy.
-    /// Resize memory in this case to the desired size.
-    const auto max_first_buffer = buffer_allocation_policy->getBufferSize();
-    if (memory.size() > max_first_buffer)
-    {
-        memory.resize(max_first_buffer);
-        WriteBuffer::set(memory.data(), memory.size());
-    }
-}
-
 void WriteBufferFromS3::setFakeBufferWhenPreFinalized()
 {
     WriteBuffer::set(fake_buffer_when_prefinalized, sizeof(fake_buffer_when_prefinalized));
diff --git a/src/IO/WriteBufferFromS3.h b/src/IO/WriteBufferFromS3.h
index b026da607c5..604f036fcb8 100644
--- a/src/IO/WriteBufferFromS3.h
+++ b/src/IO/WriteBufferFromS3.h
@@ -64,7 +64,6 @@ private:
     void reallocateFirstBuffer();
     void detachBuffer();
     void allocateBuffer();
-    void allocateFirstBuffer();
     void setFakeBufferWhenPreFinalized();
 
     S3::UploadPartRequest getUploadRequest(size_t part_number, PartData & data);