Refactor WAL log writer && IO limiter

dbms/src/Storages/Page/V3/LogFile/LogWriter.cpp

@@ -4,6 +4,7 @@
 #include <IO/WriteBufferFromFile.h>
 #include <IO/WriteHelpers.h>
 #include <Poco/Logger.h>
+#include <Storages/Page/PageUtil.h>
 #include <Storages/Page/V3/LogFile/LogFormat.h>
 #include <Storages/Page/V3/LogFile/LogWriter.h>
 #include <common/logger_useful.h>
@@ -12,50 +13,63 @@
 namespace DB::PS::V3
 {
 LogWriter::LogWriter(
-    std::unique_ptr<WriteBufferFromFileBase> && dest_,
+    String path_,
+    const FileProviderPtr & file_provider_,
     Format::LogNumberType log_number_,
     bool recycle_log_files_,
     bool manual_flush_)
-    : dest(std::move(dest_))
+    : path(path_)
+    , file_provider(file_provider_)
     , block_offset(0)
     , log_number(log_number_)
     , recycle_log_files(recycle_log_files_)
     , manual_flush(manual_flush_)
+    , write_buffer(nullptr, 0)
 {
-    // Must be `BLOCK_SIZE`, or we can not ensure the correctness of writing.
-    assert(dest->internalBuffer().size() == Format::BLOCK_SIZE);
+    log_file = file_provider->newWritableFile(
+        path,
+        EncryptionPath(path, ""),
+        false,
+        /*create_new_encryption_info_*/ false);
+
+    buffer = static_cast<char *>(alloc(buffer_size));
+    write_buffer = WriteBuffer(buffer, buffer_size);
+}
+
+void LogWriter::resetBuffer()
+{
+    write_buffer = WriteBuffer(buffer, buffer_size);
 }
 
 LogWriter::~LogWriter()
 {
-    if (dest)
-    {
-        flush();
+    log_file->fsync();
+    log_file->close();
 
-        dest->close(); // close explicitly
-    }
+    free(buffer, buffer_size);
 }
 
 size_t LogWriter::writtenBytes() const
 {
-    return dest->getMaterializedBytes();
+    return written_bytes;
 }
 
-void LogWriter::flush()
+void LogWriter::flush(const WriteLimiterPtr & write_limiter)
 {
-    dest->sync();
+    PageUtil::writeFile(log_file, written_bytes, write_buffer.buffer().begin(), write_buffer.offset(), write_limiter, false);
+    log_file->fsync();
+    written_bytes += write_buffer.offset();
+
+    // reset the write_buffer
+    resetBuffer();
 }
 
 void LogWriter::close()
 {
-    if (dest)
-    {
-        dest->close();
-        dest.reset();
-    }
+    log_file->close();
 }
 
-void LogWriter::addRecord(ReadBuffer & payload, const size_t payload_size)
+void LogWriter::addRecord(ReadBuffer & payload, const size_t payload_size, const WriteLimiterPtr & write_limiter)
 {
     // Header size varies depending on whether we are recycling or not.
     const int header_size = recycle_log_files ? Format::RECYCLABLE_HEADER_SIZE : Format::HEADER_SIZE;
@@ -64,6 +78,17 @@ void LogWriter::addRecord(ReadBuffer & payload, const size_t payload_size)
     // we still want to iterate once to emit a single zero-length record.
     bool begin = true;
     size_t payload_left = payload_size;
+
+    size_t head_sizes = ((payload_size / Format::BLOCK_SIZE) + 1) * Format::RECYCLABLE_HEADER_SIZE;
+    if (payload_size + head_sizes >= buffer_size)
+    {
+        size_t new_buff_size = payload_size + ((head_sizes / Format::BLOCK_SIZE) + 1) * Format::BLOCK_SIZE;
+
+        buffer = static_cast<char *>(realloc(buffer, buffer_size, new_buff_size));
+        buffer_size = new_buff_size;
+        resetBuffer();
+    }
+
     do
     {
         const Int64 leftover = Format::BLOCK_SIZE - block_offset;
@@ -75,7 +100,7 @@ void LogWriter::addRecord(ReadBuffer & payload, const size_t payload_size)
             {
                 // Fill the trailer with all zero
                 static constexpr char MAX_ZERO_HEADER[Format::RECYCLABLE_HEADER_SIZE]{'\x00'};
-                writeString(MAX_ZERO_HEADER, leftover, *dest);
+                writeString(MAX_ZERO_HEADER, leftover, write_buffer);
             }
             block_offset = 0;
         }
@@ -101,7 +126,9 @@ void LogWriter::addRecord(ReadBuffer & payload, const size_t payload_size)
     } while (payload.hasPendingData());
 
     if (!manual_flush)
-        dest->sync();
+    {
+        flush(write_limiter);
+    }
 }
 
 void LogWriter::emitPhysicalRecord(Format::RecordType type, ReadBuffer & payload, size_t length)
@@ -150,9 +177,9 @@ void LogWriter::emitPhysicalRecord(Format::RecordType type, ReadBuffer & payload
     // Write the checksum, header and the payload
     digest.update(payload.position(), length);
     Format::ChecksumType checksum = digest.checksum();
-    writeIntBinary(checksum, *dest);
-    writeString(header_buff.buffer().begin(), header_buff.count(), *dest);
-    writeString(payload.position(), length, *dest);
+    writeIntBinary(checksum, write_buffer);
+    writeString(header_buff.buffer().begin(), header_buff.count(), write_buffer);
+    writeString(payload.position(), length, write_buffer);
 
     block_offset += header_size + length;
 }

dbms/src/Storages/Page/V3/LogFile/LogWriter.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <Encryption/FileProvider.h>
 #include <Storages/Page/V3/LogFile/LogFormat.h>
 #include <common/types.h>
 
@@ -55,11 +56,12 @@ namespace PS::V3
  * Log number = 32bit log file number, so that we can distinguish between
  * records written by the most recent log writer vs a previous one.
  */
-class LogWriter final
+class LogWriter final : public Allocator<false>
 {
 public:
     LogWriter(
-        std::unique_ptr<WriteBufferFromFileBase> && dest_,
+        String path_,
+        const FileProviderPtr & file_provider_,
         Format::LogNumberType log_number_,
         bool recycle_log_files_,
         bool manual_flush_ = false);
@@ -69,9 +71,9 @@ class LogWriter final
 
     ~LogWriter();
 
-    void addRecord(ReadBuffer & payload, size_t payload_size);
+    void addRecord(ReadBuffer & payload, size_t payload_size, const WriteLimiterPtr & write_limiter = nullptr);
 
-    void flush();
+    void flush(const WriteLimiterPtr & write_limiter = nullptr);
 
     void close();
 
@@ -85,14 +87,26 @@ class LogWriter final
 private:
     void emitPhysicalRecord(Format::RecordType type, ReadBuffer & payload, size_t length);
 
+    void resetBuffer();
+
 private:
-    std::unique_ptr<WriteBufferFromFileBase> dest;
+    String path;
+    FileProviderPtr file_provider;
+
+    WritableFilePtr log_file;
+
     size_t block_offset; // Current offset in block
     Format::LogNumberType log_number;
     const bool recycle_log_files;
     // If true, it does not flush after each write. Instead it relies on the upper
     // layer to manually does the flush by calling ::flush()
     const bool manual_flush;
+
+    size_t written_bytes = 0;
+
+    char * buffer;
+    size_t buffer_size = Format::BLOCK_SIZE;
+    WriteBuffer write_buffer;
 };
 } // namespace PS::V3
 } // namespace DB

dbms/src/Storages/Page/V3/PageDirectory.cpp

@@ -394,13 +394,13 @@ std::set<PageId> PageDirectory::getAllPageIds()
     return page_ids;
 }
 
-void PageDirectory::apply(PageEntriesEdit && edit)
+void PageDirectory::apply(PageEntriesEdit && edit, const WriteLimiterPtr & write_limiter)
 {
     std::unique_lock write_lock(table_rw_mutex); // TODO: It is totally serialized, make it a pipeline
     UInt64 last_sequence = sequence.load();
 
     // stage 1, persisted the changes to WAL with version [seq=last_seq + 1, epoch=0]
-    wal->apply(edit, PageVersionType(last_sequence + 1, 0));
+    wal->apply(edit, PageVersionType(last_sequence + 1, 0), write_limiter);
 
     // stage 2, create entry version list for pageId. nothing need to be rollback
     std::unordered_map<PageId, std::pair<PageLock, int>> updating_locks;
@@ -496,7 +496,7 @@ void PageDirectory::apply(PageEntriesEdit && edit)
     sequence.fetch_add(1);
 }
 
-std::set<PageId> PageDirectory::gcApply(PageEntriesEdit && migrated_edit, bool need_scan_page_ids)
+std::set<PageId> PageDirectory::gcApply(PageEntriesEdit && migrated_edit, bool need_scan_page_ids, const WriteLimiterPtr & write_limiter)
 {
     {
         std::shared_lock read_lock(table_rw_mutex);
@@ -519,7 +519,7 @@ std::set<PageId> PageDirectory::gcApply(PageEntriesEdit && migrated_edit, bool n
     } // Then we should release the read lock on `table_rw_mutex`
 
     // Apply migrate edit into WAL with the increased epoch version
-    wal->apply(migrated_edit);
+    wal->apply(migrated_edit, write_limiter);
 
     if (!need_scan_page_ids)
     {
@@ -567,12 +567,12 @@ PageDirectory::getEntriesByBlobIds(const std::vector<BlobFileId> & blob_need_gc)
 }
 
 
-std::vector<PageEntriesV3> PageDirectory::gc()
+std::vector<PageEntriesV3> PageDirectory::gc(const WriteLimiterPtr & write_limiter, const ReadLimiterPtr & read_limiter)
 {
     [[maybe_unused]] bool done_anything = false;
     UInt64 lowest_seq = sequence.load();
 
-    done_anything |= wal->compactLogs();
+    done_anything |= wal->compactLogs(write_limiter, read_limiter);
 
     {
         // Cleanup released snapshots

dbms/src/Storages/Page/V3/PageDirectory.h

@@ -214,14 +214,14 @@ class PageDirectory
 
     std::set<PageId> getAllPageIds();
 
-    void apply(PageEntriesEdit && edit);
+    void apply(PageEntriesEdit && edit, const WriteLimiterPtr & write_limiter = nullptr);
 
     std::pair<std::map<BlobFileId, PageIdAndVersionedEntries>, PageSize>
     getEntriesByBlobIds(const std::vector<BlobFileId> & blob_need_gc);
 
-    std::set<PageId> gcApply(PageEntriesEdit && migrated_edit, bool need_scan_page_ids);
+    std::set<PageId> gcApply(PageEntriesEdit && migrated_edit, bool need_scan_page_ids, const WriteLimiterPtr & write_limiter = nullptr);
 
-    std::vector<PageEntriesV3> gc();
+    std::vector<PageEntriesV3> gc(const WriteLimiterPtr & write_limiter = nullptr, const ReadLimiterPtr & read_limiter = nullptr);
 
     size_t numPages() const
     {

dbms/src/Storages/Page/V3/PageStorageImpl.cpp

@@ -35,7 +35,7 @@ void PageStorageImpl::restore()
         collapsing_directory.apply(std::move(edit));
     };
     // Restore `collapsing_directory` from disk
-    auto wal = WALStore::create(callback, file_provider, delegator, /*write_limiter*/ nullptr);
+    auto wal = WALStore::create(callback, file_provider, delegator);
     // PageId max_page_id = collapsing_directory.max_applied_page_id; // TODO: return it to outer function
 
     // TODO: Now `PageDirectory::create` and `BlobStore::restore` iterate all entries in `collapsing_directory`,
@@ -78,7 +78,7 @@ void PageStorageImpl::write(DB::WriteBatch && write_batch, const WriteLimiterPtr
 
     // Persist Page data to BlobStore
     auto edit = blob_store.write(write_batch, write_limiter);
-    page_directory.apply(std::move(edit));
+    page_directory.apply(std::move(edit), write_limiter);
 }
 
 DB::PageEntry PageStorageImpl::getEntry(PageId page_id, SnapshotPtr snapshot)
@@ -178,7 +178,7 @@ void PageStorageImpl::traverse(const std::function<void(const DB::Page & page)>
     }
 }
 
-bool PageStorageImpl::gc(bool /*not_skip*/, const WriteLimiterPtr & /*write_limiter*/, const ReadLimiterPtr & /*read_limiter*/)
+bool PageStorageImpl::gc(bool /*not_skip*/, const WriteLimiterPtr & write_limiter, const ReadLimiterPtr & read_limiter)
 {
     // If another thread is running gc, just return;
     bool v = false;
@@ -192,7 +192,7 @@ bool PageStorageImpl::gc(bool /*not_skip*/, const WriteLimiterPtr & /*write_limi
 
     // 1. Do the MVCC gc, clean up expired snapshot.
     // And get the expired entries.
-    const auto & del_entries = page_directory.gc();
+    const auto & del_entries = page_directory.gc(write_limiter, read_limiter);
 
     // 2. Remove the expired entries in BlobStore.
     // It won't delete the data on the disk.
@@ -224,7 +224,7 @@ bool PageStorageImpl::gc(bool /*not_skip*/, const WriteLimiterPtr & /*write_limi
     // 5. Do the BlobStore GC
     // After BlobStore GC, these entries will be migrated to a new blob.
     // Then we should notify MVCC apply the change.
-    PageEntriesEdit gc_edit = blob_store.gc(blob_gc_info, total_page_size);
+    PageEntriesEdit gc_edit = blob_store.gc(blob_gc_info, total_page_size, write_limiter, read_limiter);
     if (gc_edit.empty())
     {
         throw Exception("Something wrong after BlobStore GC.", ErrorCodes::LOGICAL_ERROR);
@@ -237,7 +237,7 @@ bool PageStorageImpl::gc(bool /*not_skip*/, const WriteLimiterPtr & /*write_limi
     // be reset to correct state during restore. If any exception thrown, then some BlobFiles
     // will be remained as "read-only" files while entries in them are useless in actual.
     // Those BlobFiles should be cleaned during next restore.
-    const auto & page_ids = page_directory.gcApply(std::move(gc_edit), external_pages_remover != nullptr);
+    const auto & page_ids = page_directory.gcApply(std::move(gc_edit), external_pages_remover != nullptr, write_limiter);
 
     (void)page_ids;
 

dbms/src/Storages/Page/V3/WAL/WALReader.cpp

@@ -60,24 +60,25 @@ LogFilenameSet WALStoreReader::listAllFiles(
     return log_files;
 }
 
-WALStoreReaderPtr WALStoreReader::create(FileProviderPtr & provider, LogFilenameSet files)
+WALStoreReaderPtr WALStoreReader::create(FileProviderPtr & provider, LogFilenameSet files, const ReadLimiterPtr & read_limiter)
 {
-    auto reader = std::make_shared<WALStoreReader>(provider, std::move(files));
+    auto reader = std::make_shared<WALStoreReader>(provider, std::move(files), read_limiter);
     reader->openNextFile();
     return reader;
 }
 
-WALStoreReaderPtr WALStoreReader::create(FileProviderPtr & provider, PSDiskDelegatorPtr & delegator)
+WALStoreReaderPtr WALStoreReader::create(FileProviderPtr & provider, PSDiskDelegatorPtr & delegator, const ReadLimiterPtr & read_limiter)
 {
     Poco::Logger * logger = &Poco::Logger::get("WALStore");
     LogFilenameSet log_files = listAllFiles(delegator, logger);
-    return create(provider, std::move(log_files));
+    return create(provider, std::move(log_files), read_limiter);
 }
 
-WALStoreReader::WALStoreReader(FileProviderPtr & provider_, LogFilenameSet && files_)
+WALStoreReader::WALStoreReader(FileProviderPtr & provider_, LogFilenameSet && files_, const ReadLimiterPtr & read_limiter_)
     : provider(provider_)
     , files(std::move(files_))
     , next_reading_file(files.begin())
+    , read_limiter(read_limiter_)
     , logger(&Poco::Logger::get("LogReader"))
 {}
 
@@ -135,7 +136,7 @@ bool WALStoreReader::openNextFile()
             EncryptionPath{parent_path, filename},
             /*estimated_size*/ Format::BLOCK_SIZE,
             /*aio_threshold*/ 0,
-            /*read_limiter*/ nullptr,
+            /*read_limiter*/ read_limiter,
             /*buffer_size*/ Format::BLOCK_SIZE // Must be `Format::BLOCK_SIZE`
         );
         reader = std::make_unique<LogReader>(

dbms/src/Storages/Page/V3/WAL/WALReader.h

@@ -33,9 +33,9 @@ class WALStoreReader
 public:
     static LogFilenameSet listAllFiles(PSDiskDelegatorPtr & delegator, Poco::Logger * logger);
 
-    static WALStoreReaderPtr create(FileProviderPtr & provider, LogFilenameSet files);
+    static WALStoreReaderPtr create(FileProviderPtr & provider, LogFilenameSet files, const ReadLimiterPtr & read_limiter = nullptr);
 
-    static WALStoreReaderPtr create(FileProviderPtr & provider, PSDiskDelegatorPtr & delegator);
+    static WALStoreReaderPtr create(FileProviderPtr & provider, PSDiskDelegatorPtr & delegator, const ReadLimiterPtr & read_limiter = nullptr);
 
     bool remained() const;
 
@@ -56,7 +56,7 @@ class WALStoreReader
         return reader->getLogNumber();
     }
 
-    WALStoreReader(FileProviderPtr & provider_, LogFilenameSet && files_);
+    WALStoreReader(FileProviderPtr & provider_, LogFilenameSet && files_, const ReadLimiterPtr & read_limiter_ = nullptr);
 
     WALStoreReader(const WALStoreReader &) = delete;
     WALStoreReader & operator=(const WALStoreReader &) = delete;
@@ -69,6 +69,7 @@ class WALStoreReader
 
     const LogFilenameSet files;
     LogFilenameSet::const_iterator next_reading_file;
+    const ReadLimiterPtr read_limiter;
     std::unique_ptr<LogReader> reader;
     Poco::Logger * logger;
 };