file_writer.hpp

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <array>
#include <cerrno>
#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <filesystem>
#include <initializer_list>
#include <memory>
#include <mutex>
#include <span>
#include <string>
#include <system_error>
#include <thread>
#include <vector>
 
#include <fcntl.h>
#include <linux/fs.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/vfs.h>
#include <unistd.h>
 
#include <cuddl/detail/ablation.hpp>
#include <cuddl/detail/io_ring.hpp>
 
namespace cuddl::detail {
 
class file_writer {
   public:
    static constexpr size_t block_bytes = size_t{8} << 20;
    static constexpr unsigned depth = 4;
    static constexpr size_t probe_window = size_t{64} << 20;
 
    file_writer(
        std::filesystem::path const& path,
        bool durable,
        size_t queued_limit = size_t{1} << 30
    )
        : path_(path), queued_limit_(queued_limit), durable_(durable) {
        int const flags = O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC;
        fd_ = ablated("direct-io") ? -1 : ::open(path.c_str(), flags | O_DIRECT, 0644);
        direct_ = fd_ != -1;
        if (fd_ == -1) fd_ = ::open(path.c_str(), flags, 0644);
        if (fd_ == -1) fail("cannot write");
        struct statfs status{};
        // shortcut: NFS and Lustre only, the network file systems we run on; add others as met.
        if (direct_ && ::fstatfs(fd_, &status) == 0 &&
            (status.f_type == 0x6969 || status.f_type == 0x0BD00BD0)) {
            if (::fcntl(fd_, F_SETFL, ::fcntl(fd_, F_GETFL) & ~O_DIRECT) == -1) {
                ::close(fd_);
                fail("cannot write");
            }
            direct_ = false;
        }
        // Buffered writes through io_uring would run concurrently on kernel workers, and several
        // writers into one Lustre file were far slower than one.
        if (direct_ && !ablated("io-uring")) ring_ = io_ring::make(depth);
        // A drive slower than the page cache leaves a run that need not be durable waiting on
        // direct writes the cache could absorb: 159 GB to a PM9A3 took 108-158 s direct and
        // 60-97 s buffered. Such writers time both routes and keep the faster; see work().
        if (direct_ && !durable_ && !ablated("write-probe")) {
            buffered_fd_ = ::open(path.c_str(), O_WRONLY | O_CLOEXEC);
            if (buffered_fd_ == -1) fail("cannot write");
            // btrfs compresses page-cache writes during writeback, never direct ones; with zstd
            // a buffered 120 GB run took 58.5 s compressed and 45.8 s not. Other file systems
            // refuse or ignore the flag.
            int attributes = 0;
            if (::ioctl(fd_, FS_IOC_GETFLAGS, &attributes) == 0) {
                attributes |= FS_NOCOMP_FL;
                static_cast<void>(::ioctl(fd_, FS_IOC_SETFLAGS, &attributes));
            }
        }
        thread_ = std::thread([this] { work(); });
    }
    file_writer(file_writer const&) = delete;
    file_writer& operator=(file_writer const&) = delete;
    ~file_writer() {
        stop();
        if (fd_ != -1) ::close(fd_);
        if (buffered_fd_ != -1) ::close(buffered_fd_);
    }
 
    std::span<char> reserve(size_t bytes) {
        std::unique_lock lock(append_mutex_);
        auto const total = carry_bytes_ + bytes;
        current_ = take_extent((total + 4095) / 4096 * 4096);
        std::memcpy(current_->data.get(), carry_.data(), carry_bytes_);
        reserved_ = bytes;
        reservation_ = std::move(lock);
        return {current_->data.get() + carry_bytes_, bytes};
    }
 
    void commit() {
        queue_current(false);
        reservation_.unlock();
    }
 
    void append(std::initializer_list<std::span<char const>> chunks) {
        size_t bytes = 0;
        for (auto const& chunk : chunks) bytes += chunk.size();
        auto* out = reserve(bytes).data();
        for (auto const& chunk : chunks) {
            if (!chunk.empty()) std::memcpy(out, chunk.data(), chunk.size());
            out += chunk.size();
        }
        commit();
    }
 
    void append(std::span<char const> bytes) {
        append({bytes});
    }
 
    void finish() {
        uint64_t size = 0;
        {
            reserve(0);
            size = offset_ + carry_bytes_;
            queue_current(true);
            reservation_.unlock();
        }
        {
            std::unique_lock lock(mutex_);
            space_.wait(lock, [&] { return failed_ != 0 || (queue_.empty() && in_flight_ == 0); });
        }
        stop();
        int error = failed_;
        if (error == 0 && direct_ && ::ftruncate(fd_, static_cast<off_t>(size)) != 0) error = errno;
        if (error == 0 && durable_ && ::fsync(fd_) != 0) error = errno;
        if (::close(fd_) != 0 && error == 0) error = errno;
        fd_ = -1;
        if (buffered_fd_ != -1 && ::close(buffered_fd_) != 0 && error == 0) error = errno;
        buffered_fd_ = -1;
        if (error != 0) fail("cannot write", error);
    }
 
    [[nodiscard]] bool direct() const noexcept {
        return direct_;
    }
 
    [[nodiscard]] bool asynchronous() const noexcept {
        return ring_ != nullptr;
    }
 
   private:
    struct free_deleter {
        void operator()(char* data) const noexcept {
            std::free(data);
        }
    };
 
    struct extent {
        std::unique_ptr<char, free_deleter> data;
        size_t capacity = 0;
        unsigned pending = 0;  // queued pieces not yet written
    };
 
    struct queued_block {
        extent* source = nullptr;
        char const* data = nullptr;
        size_t bytes = 0;
        uint64_t offset = 0;
    };
 
    [[noreturn]] void fail(char const* what, int error = errno) const {
        throw std::system_error(
            error, std::generic_category(), std::string(what) + " " + path_.string()
        );
    }
 
    extent* take_extent(size_t bytes) {
        std::unique_lock lock(mutex_);
        space_.wait(lock, [&] {
            return failed_ != 0 || extents_.size() - free_.size() < 2 ||
                   queued_ + bytes <= queued_limit_;
        });
        if (failed_ != 0) fail("cannot write", failed_);
        auto const fits =
            std::ranges::find_if(free_, [&](extent* e) { return e->capacity >= bytes; });
        extent* chosen = nullptr;
        if (fits != free_.end()) {
            chosen = *fits;
            free_.erase(fits);
        } else {
            // Too small for anything to come, so its memory goes to a larger extent.
            if (!free_.empty()) {
                auto* const smallest = free_.front();
                free_.erase(free_.begin());
                std::erase_if(extents_, [&](auto const& e) { return e.get() == smallest; });
            }
            auto fresh = std::make_unique<extent>();
            // O_DIRECT needs memory aligned to the logical sector, 4096 bytes at most; madvise
            // needs whole pages, which are larger on some systems.
            static size_t const alignment =
                std::max<size_t>(4096, static_cast<size_t>(::sysconf(_SC_PAGESIZE)));
            fresh->capacity =
                (std::max<size_t>(bytes, 1) + block_bytes - 1) / block_bytes * block_bytes;
            fresh->data.reset(static_cast<char*>(std::aligned_alloc(alignment, fresh->capacity)));
            if (!fresh->data) throw std::bad_alloc();
            // glibc backs large allocations with transparent huge pages, 512 MiB on a 64 KiB-page
            // kernel. With 72 threads filling them, a GH200 CLI search writing 120 GB took 15 s
            // of kernel time and 1.38 s against 1.15 s on small pages. The advice is only a hint.
            static_cast<void>(::madvise(fresh->data.get(), fresh->capacity, MADV_NOHUGEPAGE));
            chosen = extents_.emplace_back(std::move(fresh)).get();
        }
        queued_ += chosen->capacity;
        return chosen;
    }
 
    void queue_current(bool final) {
        auto const total = carry_bytes_ + reserved_;
        size_t written = total;
        if (direct_) written = final ? (total + 4095) / 4096 * 4096 : total / 4096 * 4096;
        char* const data = current_->data.get();
        std::memset(data + total, 0, written > total ? written - total : 0);
        carry_bytes_ = total - std::min(total, written);
        std::memcpy(carry_.data(), data + written, carry_bytes_);
        {
            std::lock_guard lock(mutex_);
            for (size_t at = 0; at < written; at += block_bytes) {
                queue_.push_back(
                    {current_, data + at, std::min(block_bytes, written - at), offset_ + at}
                );
                ++current_->pending;
            }
            if (current_->pending == 0) release(current_);
        }
        offset_ += written;
        current_ = nullptr;
        reserved_ = 0;
        ready_.notify_all();
    }
 
    void release(extent* done) {
        queued_ -= done->capacity;
        auto const at = std::ranges::lower_bound(free_, done->capacity, {}, &extent::capacity);
        free_.insert(at, done);
    }
 
    void stop() {
        {
            std::lock_guard lock(mutex_);
            stopping_ = true;
        }
        ready_.notify_all();
        if (thread_.joinable()) thread_.join();
    }
 
    static int write_all(int fd, queued_block const& piece) {
        size_t done = 0;
        while (done < piece.bytes) {
            auto const wrote = ::pwrite(
                fd, piece.data + done, piece.bytes - done, static_cast<off_t>(piece.offset + done)
            );
            if (wrote < 0) return errno;
            if (wrote == 0) return EIO;
            done += static_cast<size_t>(wrote);
        }
        return 0;
    }
 
    using clock = std::chrono::steady_clock;
 
    static double seconds_since(clock::time_point start) {
        return std::chrono::duration<double>(clock::now() - start).count();
    }
 
    void retire(queued_block const& written, int error) {
        {
            std::lock_guard lock(mutex_);
            if (error != 0 && failed_ == 0) failed_ = error;
            if (--written.source->pending == 0) release(written.source);
        }
        space_.notify_all();
    }
 
    void work() {
        std::vector<queued_block> slots(ring_ ? depth : 1U);
        std::vector<unsigned> idle;
        for (unsigned slot = 0; slot < slots.size(); ++slot) idle.push_back(slot);
        auto const window = std::min(probe_window, queued_limit_);
        bool buffered = false, settled = buffered_fd_ == -1, tried = false;
        uint64_t submitted = 0, window_bytes = 0;
        double window_seconds = 0, direct_rate = 0;
        while (true) {
            if (buffered) {
                queued_block next;
                {
                    std::unique_lock lock(mutex_);
                    ready_.wait(lock, [&] { return stopping_ || !queue_.empty(); });
                    if (queue_.empty()) return;
                    next = queue_.front();
                    queue_.pop_front();
                    ++in_flight_;
                }
                auto const start = clock::now();
                int const error = write_all(buffered_fd_, next);
                window_seconds += seconds_since(start);
                window_bytes += next.bytes;
                {
                    std::lock_guard lock(mutex_);
                    --in_flight_;
                }
                retire(next, error);
                if (window_bytes >= (tried ? 16 : 1) * window) {
                    auto const pace = tried ? 0.75 * direct_rate : direct_rate;
                    buffered = static_cast<double>(window_bytes) >= pace * window_seconds;
                    settled = !buffered;
                    tried = true;
                    window_bytes = 0;
                    window_seconds = 0;
                }
                continue;
            }
            {
                std::unique_lock lock(mutex_);
                // With writes in flight, collect their completions before waiting for more.
                if (idle.size() == slots.size()) {
                    if (!settled && submitted >= window) {
                        direct_rate = static_cast<double>(window_bytes) / window_seconds;
                        buffered = true;
                        window_bytes = 0;
                        window_seconds = 0;
                        continue;
                    }
                    ready_.wait(lock, [&] { return stopping_ || !queue_.empty(); });
                    if (queue_.empty()) return;
                }
                for (; !idle.empty() && !queue_.empty() && (settled || submitted < window);
                     idle.pop_back()) {
                    slots[idle.back()] = queue_.front();
                    queue_.pop_front();
                    ++in_flight_;
                    submitted += slots[idle.back()].bytes;
                    if (ring_) {
                        auto const& next = slots[idle.back()];
                        ring_->write(
                            fd_,
                            next.data,
                            static_cast<uint32_t>(next.bytes),
                            next.offset,
                            idle.back()
                        );
                    }
                }
            }
            if (idle.size() == slots.size()) continue;
            auto const start = clock::now();
            if (!ring_) {
                int const error = write_all(fd_, slots[0]);
                window_seconds += seconds_since(start);
                window_bytes += slots[0].bytes;
                finish_slot(slots, idle, 0, error);
                continue;
            }
            int const error = ring_->submit_and_wait(1, [&](uint64_t slot, int result) {
                auto const expected = static_cast<int64_t>(slots[slot].bytes);
                window_bytes += slots[slot].bytes;
                // A short direct write means the device ran out of space or failed.
                finish_slot(
                    slots,
                    idle,
                    static_cast<unsigned>(slot),
                    result < 0 ? -result : (result == expected ? 0 : EIO)
                );
            });
            window_seconds += seconds_since(start);
            if (error != 0) {
                // The kernel may still read the extents in flight, so they are never freed.
                for (auto& extent : extents_) static_cast<void>(extent->data.release());
                std::lock_guard lock(mutex_);
                failed_ = error;
                space_.notify_all();
                return;
            }
        }
    }
 
    void finish_slot(
        std::vector<queued_block>& slots,
        std::vector<unsigned>& idle,
        unsigned slot,
        int error
    ) {
        {
            std::lock_guard lock(mutex_);
            --in_flight_;
        }
        retire(slots[slot], error);
        idle.push_back(slot);
    }
 
    std::filesystem::path path_;
    size_t const queued_limit_;
    bool const durable_;
    int fd_ = -1;
    int buffered_fd_ = -1;
    bool direct_ = false;
    std::unique_ptr<io_ring> ring_;
 
    // Producer state, guarded by append_mutex_, which a reservation holds until its commit.
    std::mutex append_mutex_;
    std::unique_lock<std::mutex> reservation_;
    extent* current_ = nullptr;
    size_t reserved_ = 0;
    uint64_t offset_ = 0;
    std::array<char, 4096> carry_{};
    size_t carry_bytes_ = 0;
 
    std::mutex mutex_;
    std::condition_variable ready_, space_;
    std::deque<queued_block> queue_;
    std::vector<std::unique_ptr<extent>> extents_;
    std::vector<extent*> free_;
    size_t queued_ = 0;
    unsigned in_flight_ = 0;
    int failed_ = 0;
    bool stopping_ = false;
    std::thread thread_;
};
 
}  // namespace cuddl::detail