BufferedWriter.hpp

Go to the documentation of this file.

#ifndef BYTEME_BUFFERED_WRITER_HPP
#define BYTEME_BUFFERED_WRITER_HPP
 
#include <thread>
#include <condition_variable>
#include <mutex>
#include <vector>
#include <algorithm>
#include <exception>
#include <type_traits>
#include <memory>
#include <cstddef>
#include <string>
 
#include "sanisizer/sanisizer.hpp"
 
#include "Writer.hpp"
#include "utils.hpp"
 
namespace byteme {
 
template<typename Type_, class WriterPointer_ = std::unique_ptr<Writer> >
class BufferedWriter {
public:
    BufferedWriter(std::size_t buffer_size) : my_buffer(sanisizer::cast<I<decltype(my_buffer.size())> >(buffer_size)) {
        if (buffer_size == 0) {
            throw std::runtime_error("buffer size must be positive");
        }
    }
 
    BufferedWriter(BufferedWriter&&) = delete;
    BufferedWriter(const BufferedWriter&) = delete; // not copyable.
    BufferedWriter& operator=(BufferedWriter&&) = delete;
    BufferedWriter& operator=(const BufferedWriter&) = delete; // not copyable.
 
    virtual ~BufferedWriter() = default;
protected:
    virtual void flush_async() = 0;
 
    virtual void flush_sync(const Type_*, std::size_t) = 0;
 
    std::vector<Type_>& get_buffer() {
        return my_buffer;
    }
private:
    std::vector<Type_> my_buffer;
    std::size_t my_current = 0;
 
    // Standard guarantees this to be at least 64 bits, which is more than that of size_t.
    // We don't use uint64_t because that might not be defined by the implementation.
    unsigned long long my_overall = 0;
 
public:
    unsigned long long number() {
        return my_overall + my_current;
    }
 
public:
    void write(Type_ input) {
        my_buffer[my_current] = input;
        ++my_current;
 
        const auto buffer_size = my_buffer.size();
        if (my_current == buffer_size) {
            flush_async();
            my_current = 0;
            my_overall += buffer_size;
        }
    }
 
    void write(const Type_* input, std::size_t number) {
        const auto buffer_size = my_buffer.size();
 
        // Seeing if we fill up the rest of our buffer or not.
        const std::size_t remaining = buffer_size - my_current;
        if (number < remaining) {
            std::copy_n(input, number, my_buffer.data() + my_current); 
            my_current += number;
            return;
        }
 
        std::copy_n(input, remaining, my_buffer.data() + my_current); 
        number -= remaining;
        input += remaining;
 
        if (number == 0) {
            flush_async();
            my_overall += buffer_size;
            my_current = 0;
            return;
        }
 
        flush_sync(my_buffer.data(), buffer_size);
        my_overall += buffer_size;
 
        // Directly filling the output array, bypassing our buffer.
        while (number >= buffer_size) {
            flush_sync(input, buffer_size);
            my_overall += buffer_size;
            input += buffer_size;
            number -= buffer_size;
        } 
 
        // Filling the cache and copying the remainder into the output array.
        std::copy_n(input, number, my_buffer.data());
        my_current = number;
    }
 
    void write(const char* string) {
        write(reinterpret_cast<const Type_*>(string), std::char_traits<char>::length(string));
    }
 
    void write(const std::string& string) {
        write(reinterpret_cast<const Type_*>(string.c_str()), string.size());
    }
 
public:
    void flush() {
        flush_sync(my_buffer.data(), my_current); // Flush whatever's left.
        my_overall += my_current;
        my_current = 0;
    }
 
    virtual void finish() = 0;
};
 
template<typename Type_, class Pointer_ = std::unique_ptr<Writer> >
class SerialBufferedWriter final : public BufferedWriter<Type_> {
public:
    SerialBufferedWriter(Pointer_ writer, std::size_t buffer_size) : BufferedWriter<Type_>(buffer_size), my_writer(std::move(writer)) {}
 
    ~SerialBufferedWriter() {
        this->flush();
    }
private:
    Pointer_ my_writer;
 
protected:
    void flush_async() {
        auto& buffer = this->get_buffer();
        my_writer->write(reinterpret_cast<unsigned char*>(buffer.data()), buffer.size());
    }
 
    void flush_sync(const Type_* ptr, std::size_t num) {
        my_writer->write(reinterpret_cast<const unsigned char*>(ptr), num);
    }
public:
    void finish() {
        this->flush();
        my_writer->finish();
    }
};
 
template<typename Type_, class Pointer_ = std::unique_ptr<Writer> >
class ParallelBufferedWriter final : public BufferedWriter<Type_> {
public:
    ParallelBufferedWriter(Pointer_ writer, std::size_t buffer_size) : 
        BufferedWriter<Type_>(buffer_size),
        my_writer(std::move(writer)),
        my_buffer_worker(sanisizer::cast<I<decltype(my_buffer_worker.size())> >(buffer_size))
    {
        my_thread = std::thread([&]() { thread_loop(); }); // set up thread before initializing.
    }
 
    ~ParallelBufferedWriter() {
        this->flush();
 
        std::unique_lock lck(my_mut);
        my_kill = true;
        my_ready_input = true;
        lck.unlock(); // releasing the lock so that the notified thread doesn't immediately block.
        my_cv.notify_one();
        my_thread.join();
    }
private:
    Pointer_ my_writer;
    std::vector<Type_> my_buffer_worker;
    std::size_t my_to_write = 0;
 
private:
    std::thread my_thread;
    std::exception_ptr my_thread_err = nullptr;
    std::mutex my_mut;
    std::condition_variable my_cv;
 
    bool my_ready_input = false, my_ready_output = false;
    bool my_worker_active = false;
    bool my_kill = false;
 
    void thread_loop() {
        while (1) {
            std::unique_lock lck(my_mut);
            my_cv.wait(lck, [&]() { return my_ready_input; });
            my_ready_input = false;
 
            if (my_kill) { // Handle an explicit kill signal from the destructor.
                break;
            }
 
            try {
                my_writer->write(reinterpret_cast<const unsigned char*>(my_buffer_worker.data()), my_to_write);
            } catch (...) {
                my_thread_err = std::current_exception();
            }
 
            my_ready_output = true;
            lck.unlock();
            my_cv.notify_one();
        }
    }
 
protected:
    void flush_async() {
        auto& buffer = this->get_buffer();
        const auto num = buffer.size();
 
        if (my_worker_active) {
            // If the worker is active, it's probably already gotten started so we just wait for it to finish.
            // Then we swap the results with the main buffer and submit a new job to the worker.
            std::unique_lock lck(my_mut);
            my_cv.wait(lck, [&]() { return my_ready_output; });
            my_ready_output = false;
 
            if (my_thread_err) {
                std::rethrow_exception(my_thread_err);
            }
 
            buffer.swap(my_buffer_worker);
            my_to_write = num;
 
            my_ready_input = true;
            lck.unlock();
            my_cv.notify_one();
 
        } else {
            // If the worker is not active, we can just submit directly.
            std::unique_lock lck(my_mut);
            buffer.swap(my_buffer_worker);
            my_to_write = num;
 
            my_ready_input = true;
            lck.unlock();
            my_cv.notify_one();
            my_worker_active = true;
        }
    }
 
    void flush_sync(const Type_* ptr, std::size_t num) {
        if (my_worker_active) {
            // If the worker is active, we wait for it to finish and then write from the supplied pointer.
            // We do not submit a new job, based on the loop in the BufferedWriter::write() overload:
            // 
            // - We'll probably want to call this flush() overload again.
            //   On subsequent calls, we'll want to directly write from the supplied pointer to the Writer to cut out a copy.
            // - If we didn't call this overload again, we'd probably call the other refill() overload immediately, 
            //   So if we sent a new job to the worker, we'd end up immediately blocking on it to submit the next write.
            //   So we might as well skip that communication overhead and write directly on the main thread.
            std::unique_lock lck(my_mut);
            my_cv.wait(lck, [&]() { return my_ready_output; });
            my_ready_output = false;
            if (my_thread_err) {
                std::rethrow_exception(my_thread_err);
            }
 
            my_writer->write(reinterpret_cast<const unsigned char*>(ptr), num);
            my_worker_active = false;
 
        } else {
            // If the worker isn't active, we can directly write from the supplied pointer, avoiding some communication overhead.
            my_writer->write(reinterpret_cast<const unsigned char*>(ptr), num);
        }
    }
public:
    void finish() {
        this->flush();
        my_writer->finish();
    }
};
 
}
 
#endif