fluxen.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <atomic>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <filesystem>
#include <functional>
#include <mutex>
#include <optional>
#include <shared_mutex>
#include <span>
#include <stdexcept>
#include <string>
#include <string_view>
#include <type_traits>
#include <unordered_map>
#include <vector>
 
#ifdef _WIN32
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#else
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#endif
 
namespace fluxen {
 
// --- public types ---

using Bytes = std::span<const std::byte>;

enum TxResult : uint8_t { commit, rollback };

struct io_error : public std::runtime_error {
  using std::runtime_error::runtime_error;
};
 
struct corrupt_error : public std::runtime_error {
  using std::runtime_error::runtime_error;
};
 
struct poisoned_error : public io_error {
  using io_error::io_error;
};
 
struct key_error : public std::runtime_error {
  using std::runtime_error::runtime_error;
};
 
// --- internal ---

namespace detail {
 
inline constexpr uint8_t MAGIC[8]   = {'F', 'L', 'U', 'X', 'E', 'N', '0', '1'};
inline constexpr uint8_t FLAG_LIVE  = 0x00;
inline constexpr uint8_t FLAG_TOMB  = 0x01;
inline constexpr uint8_t MAX_KEY    = 255;
inline constexpr size_t HEADER_SIZE = 6; // on-disk entry header size
 
/* On-disk entry header (6 bytes on disk) */
struct EntryHeader {
  uint8_t flags;
  uint8_t key_len;
  uint32_t val_len;
};
 
/* Serialize header into 6-byte buffer */
inline void encode_header(uint8_t out[HEADER_SIZE],
                          const EntryHeader &h) noexcept {
  out[0] = h.flags;
  out[1] = h.key_len;
  out[2] = static_cast<uint8_t>(h.val_len & 0xFFU);
  out[3] = static_cast<uint8_t>((h.val_len >> 8) & 0xFFU);
  out[4] = static_cast<uint8_t>((h.val_len >> 16) & 0xFFU);
  out[5] = static_cast<uint8_t>((h.val_len >> 24) & 0xFFU);
}
 
/* Deserialise header from a 6-byte buffer */
inline EntryHeader decode_header(const uint8_t in[HEADER_SIZE]) noexcept {
  return {
      .flags   = in[0],
      .key_len = in[1],
      .val_len = static_cast<uint32_t>(in[2]) |
                 static_cast<uint32_t>(in[3]) << 8 |
                 static_cast<uint32_t>(in[4]) << 16 |
                 static_cast<uint32_t>(in[5]) << 24,
  };
}
 
/* In-memory index entry (points into the mmap'd file) */
struct IndexEntry {
  size_t val_offset; // byte offset of value data in file
  uint32_t val_len;
};
 
struct StringHash {
  using is_transparent = void;
 
  size_t operator()(std::string_view sv) const noexcept {
    return std::hash<std::string_view>{}(sv);
  }
};
 
#ifdef _WIN32
std::wstring utf8_to_wstring(std::string_view utf8) {
  if (utf8.empty()) {
    return {};
  }
 
  int len = MultiByteToWideChar(CP_UTF8, 0, utf8.data(),
                                static_cast<int>(utf8.size()), nullptr, 0);
  if (len <= 0) {
    return {};
  }
 
  std::wstring wide(len, L'\0');
  MultiByteToWideChar(CP_UTF8, 0, utf8.data(), static_cast<int>(utf8.size()),
                      wide.data(), len);
 
  return wide;
}
#endif
 
using IndexMap =
    std::unordered_map<std::string, IndexEntry, StringHash, std::equal_to<>>;
 
/* Cross-platform mmap wrapper */
class MappedFile {
private:
#ifdef _WIN32
  HANDLE file_ = INVALID_HANDLE_VALUE;
  HANDLE map_  = nullptr;
#else
  int fd_ = -1;
#endif
  uint8_t *ptr_     = nullptr;
  size_t size_      = 0;
  size_t file_size_ = 0;
  std::atomic<bool> dirty_{false};
  std::filesystem::path path_;
 
public:
  MappedFile() = default;
  ~MappedFile() { close(); }
 
  MappedFile(const MappedFile &)            = delete;
  MappedFile &operator=(const MappedFile &) = delete;
  MappedFile(const MappedFile &&)           = delete;
  MappedFile &operator=(MappedFile &&)      = delete;
 
  bool open(std::string_view path) {
#ifdef _WIN32
    path_ = utf8_to_wstring(path);
    file_ = CreateFileW(path_.c_str(), GENERIC_READ | GENERIC_WRITE, 0, nullptr,
                        OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr);
    if (file_ == INVALID_HANDLE_VALUE) {
      return false;
    }
    SetFilePointer(file_, 0, nullptr, FILE_END);
#else
    path_ = path;
    fd_   = ::open(path_.c_str(), O_RDWR | O_CREAT | O_APPEND, 0644);
    if (fd_ < 0) {
      return false;
    }
#endif
    return remap();
  }
 
  void close() {
    unmap();
#ifdef _WIN32
    if (file_ != INVALID_HANDLE_VALUE) {
      CloseHandle(file_);
      file_ = INVALID_HANDLE_VALUE;
    }
#else
    if (fd_ >= 0) {
      ::close(fd_);
      fd_ = -1;
    }
#endif
  }
 
  bool remap() {
    unmap();
    size_      = file_size();
    file_size_ = size_;
    if (size_ == 0) {
      dirty_.store(false, std::memory_order_release);
      return true;
    }
 
#ifdef _WIN32
    map_ = CreateFileMappingW(file_, nullptr, PAGE_READWRITE, 0, 0, nullptr);
 
    if (!map_) {
      dirty_.store(false, std::memory_order_release);
      return false;
    }
 
    ptr_ = static_cast<uint8_t *>(
        MapViewOfFile(map_, FILE_MAP_ALL_ACCESS, 0, 0, size_));
 
    if (!ptr_) {
      CloseHandle(map_);
      map_ = nullptr;
      dirty_.store(false, std::memory_order_release);
      return false;
    }
#else
    ptr_ = static_cast<uint8_t *>(
        ::mmap(nullptr, size_, PROT_READ | PROT_WRITE, MAP_SHARED, fd_, 0));
 
    if (ptr_ == MAP_FAILED) {
      ptr_ = nullptr;
      dirty_.store(false, std::memory_order_release);
      return false;
    }
#endif
    dirty_.store(false, std::memory_order_release);
    return true;
  }
 
  bool append(const void *data, size_t len) {
    if (len == 0) {
      return true;
    }
#ifdef _WIN32
    DWORD written = 0;
    if (!WriteFile(file_, data, static_cast<DWORD>(len), &written, nullptr) ||
        written != static_cast<DWORD>(len)) {
      return false;
    }
#else
    if (::write(fd_, data, len) != static_cast<ssize_t>(len)) {
      return false;
    }
#endif
    dirty_.store(true, std::memory_order_relaxed);
    file_size_ += len;
    return true;
  }

  [[nodiscard]] bool sync() {
#ifdef _WIN32
    return FlushFileBuffers(file_) != 0;
#else
    return ::fsync(fd_) == 0;
#endif
  }

  [[nodiscard]] bool truncate(size_t new_size) {
#ifdef _WIN32
    unmap();
    LARGE_INTEGER li{};
    li.QuadPart = static_cast<LONGLONG>(new_size);
    if (!SetFilePointerEx(file_, li, nullptr, FILE_BEGIN)) {
      return false;
    }
    if (!SetEndOfFile(file_)) {
      SetFilePointer(file_, 0, nullptr, FILE_END);
      return false;
    }
    SetFilePointer(file_, 0, nullptr, FILE_END);
#else
    if (::ftruncate(fd_, static_cast<off_t>(new_size)) != 0) {
      return false;
    }
#endif
    file_size_ = new_size;
    dirty_.store(true, std::memory_order_release);
    return true;
  }

  bool rewrite(const std::vector<uint8_t> &data) {
    std::filesystem::path tmp_path = path_;
    tmp_path += ".tmp";
 
#ifdef _WIN32
    HANDLE tmp = CreateFileW(tmp_path.c_str(), GENERIC_WRITE, 0, nullptr,
                             CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, nullptr);
    if (tmp == INVALID_HANDLE_VALUE) {
      return false;
    }
    DWORD written = 0;
    const bool write_ok =
        WriteFile(tmp, data.data(), static_cast<DWORD>(data.size()), &written,
                  nullptr) &&
        written == static_cast<DWORD>(data.size()) &&
        FlushFileBuffers(tmp) != 0;
 
    CloseHandle(tmp);
 
    if (!write_ok) {
      DeleteFileW(tmp_path.c_str());
      return false;
    }
#else
    const int tmp_fd =
        ::open(tmp_path.c_str(), O_RDWR | O_CREAT | O_TRUNC, 0644);
    if (tmp_fd < 0) {
      return false;
    }
    const bool write_ok = ::write(tmp_fd, data.data(), data.size()) ==
                              static_cast<ssize_t>(data.size()) &&
                          ::fsync(tmp_fd) == 0;
 
    ::close(tmp_fd);
 
    if (!write_ok) {
      ::unlink(tmp_path.c_str());
      return false;
    }
#endif
    unmap();
#ifdef _WIN32
    if (file_ != INVALID_HANDLE_VALUE) {
      CloseHandle(file_);
      file_ = INVALID_HANDLE_VALUE;
    }
#else
    if (fd_ >= 0) {
      ::close(fd_);
      fd_ = -1;
    }
#endif
 
#ifdef _WIN32
    if (!ReplaceFileW(path_.c_str(), tmp_path.c_str(), nullptr,
                      REPLACEFILE_IGNORE_MERGE_ERRORS, nullptr, nullptr)) {
 
      DeleteFileW(tmp_path.c_str());
 
      file_ =
          CreateFileW(path_.c_str(), GENERIC_READ | GENERIC_WRITE, 0, nullptr,
                      OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
 
      if (file_ == INVALID_HANDLE_VALUE) {
        throw io_error("fluxen: failed to reopen database file after replace");
      }
 
      SetFilePointer(file_, 0, nullptr, FILE_END);
      if (!remap()) {
        throw io_error("fluxen: remap failed after replace failure");
      }
      return false;
    }
#else
    if (::rename(tmp_path.c_str(), path_.c_str()) != 0) {
      ::unlink(tmp_path.c_str());
      fd_ = ::open(path_.c_str(), O_RDWR | O_APPEND, 0644);
      if (fd_ < 0) {
        throw io_error("fluxen: failed to reopen database file after rename");
      }
      if (!remap()) {
        throw io_error("fluxen: remap failed after rename failure");
      }
      return false;
    }
#endif
 
#ifdef _WIN32
    file_ = CreateFileW(path_.c_str(), GENERIC_READ | GENERIC_WRITE, 0, nullptr,
                        OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
    if (file_ == INVALID_HANDLE_VALUE) {
      throw io_error(
          "fluxen: failed to reopen database file after successful replace");
    }
    SetFilePointer(file_, 0, nullptr, FILE_END);
#else
    fd_ = ::open(path_.c_str(), O_RDWR | O_APPEND, 0644);
    if (fd_ < 0) {
      throw io_error(
          "fluxen: failed to reopen database file after successful rename");
    }
#endif
    if (!remap()) {
      throw io_error("fluxen: remap failed after successful rename");
    }
    return true;
  }

  [[nodiscard]] const uint8_t *ptr() const noexcept { return ptr_; }

  [[nodiscard]] bool is_dirty() const noexcept {
    return dirty_.load(std::memory_order_acquire);
  }
 
  [[nodiscard]] size_t size() const noexcept { return file_size_; }
 
private:
  void unmap() {
    if (!ptr_)
      return;
#ifdef _WIN32
    UnmapViewOfFile(ptr_);
    if (map_) {
      CloseHandle(map_);
      map_ = nullptr;
    }
#else
    ::munmap(ptr_, size_);
#endif
    ptr_  = nullptr;
    size_ = 0;
  }
 
  [[nodiscard]] size_t file_size() const noexcept {
#ifdef _WIN32
    LARGE_INTEGER sz{};
    GetFileSizeEx(file_, &sz);
    return static_cast<size_t>(sz.QuadPart);
#else
    struct stat st{};
    ::fstat(fd_, &st);
    return static_cast<size_t>(st.st_size);
#endif
  }
};
} // namespace detail
 
// --- transaction ---
 
class DB; // forward declaration

class Tx {
private:
  friend class DB;
 
  struct Op {
    std::string key;
    std::vector<uint8_t> val;
    bool is_delete;
  };
  std::vector<Op> ops_;
 
public:
  void put(std::string_view key, std::string_view value) {
    if (key.empty() || key.size() > detail::MAX_KEY) {
      throw key_error("fluxen: key must be between 1 and 255 bytes");
    }
    ops_.push_back({.key = std::string(key),
                    .val = std::vector<uint8_t>(value.begin(), value.end()),
                    .is_delete = false});
  }

  template <typename T>
    requires(std::is_trivially_copyable_v<T> &&
             !std::is_convertible_v<T, std::string_view>)
  void put(std::string_view key, const T &value) {
    if (key.empty() || key.size() > detail::MAX_KEY) {
      throw key_error("fluxen: key must be between 1 and 255 bytes");
    }
    const auto *p = reinterpret_cast<const uint8_t *>(&value);
    ops_.push_back({.key       = std::string(key),
                    .val       = std::vector<uint8_t>(p, p + sizeof(T)),
                    .is_delete = false});
  }

  void remove(std::string_view key) {
    if (key.empty() || key.size() > detail::MAX_KEY) {
      throw key_error("fluxen: key must be between 1 and 255 bytes");
    }
    ops_.push_back({.key = std::string(key), .val = {}, .is_delete = true});
  }
};
 
// --- DB ---

class DB {
private:
  detail::IndexMap index_;
  mutable detail::MappedFile file_;
  mutable std::shared_mutex mu_;
  mutable std::mutex sync_mutex_;
  bool poisoned_ = false;
 
  void check_poisoned() const {
    if (poisoned_) {
      throw poisoned_error(
          "fluxen: database is poisoned due to an unrecoverable I/O error");
    }
  }
 
public:
  explicit DB(std::string_view path) {
    if (!file_.open(path)) {
      throw io_error("fluxen: failed to open '" + std::string(path) + "'");
    }
 
    if (file_.size() == 0) {
      init_file();
    } else {
      load_index();
    }
  }

  ~DB() = default;
  DB(const DB &)            = delete;
  DB &operator=(const DB &) = delete;
 
  // --- WRITE ---

  void put(std::string_view key, std::string_view value) {
    check_poisoned();
    std::unique_lock lock(mu_);
    append_entry(key, reinterpret_cast<const uint8_t *>(value.data()),
                 static_cast<uint32_t>(value.size()), false);
  }

  template <typename T>
    requires(std::is_trivially_copyable_v<T> &&
             !std::is_convertible_v<T, std::string_view>)
  void put(std::string_view key, const T &value) {
    check_poisoned();
    std::unique_lock lock(mu_);
    append_entry(key, reinterpret_cast<const uint8_t *>(&value),
                 static_cast<uint32_t>(sizeof(T)), false);
  }
 
  // --- READ ---

  template <typename T = std::string>
  std::optional<T> get(std::string_view key) const {
    check_poisoned();
    std::shared_lock lock(mu_);
    ensure_mapped();
 
    auto it = index_.find(key);
    if (it == index_.end()) {
      return std::nullopt;
    }
 
    const auto &entry = it->second;
 
    if constexpr (std::is_same_v<T, std::string>) {
      const auto *ptr = file_.ptr() + entry.val_offset;
      return std::string(reinterpret_cast<const char *>(ptr), entry.val_len);
    } else {
      static_assert(std::is_trivially_copyable_v<T>,
                    "fluxen: T must be trivially copyable");
      if (entry.val_len != static_cast<uint32_t>(sizeof(T))) {
        return std::nullopt;
      }
      T result;
      std::memcpy(&result, file_.ptr() + entry.val_offset, sizeof(T));
      return result;
    }
  }
 
  // --- DELETE ---

  void remove(std::string_view key) {
    check_poisoned();
    std::unique_lock lock(mu_);
    append_entry(key, nullptr, 0, true);
  }
 
  // --- QUERY ---

  [[nodiscard]] bool has(std::string_view key) const {
    check_poisoned();
    std::shared_lock lock(mu_);
    ensure_mapped();
    return index_.contains(key);
  }

  void each(const std::function<void(std::string_view, Bytes)> &fn) const {
    check_poisoned();
    std::shared_lock lock(mu_);
    ensure_mapped();
    for (const auto &[key, entry] : index_) {
      const auto *ptr =
          reinterpret_cast<const std::byte *>(file_.ptr() + entry.val_offset);
      fn(key, Bytes{ptr, entry.val_len});
    }
  }

  void prefix(std::string_view pfx,
              const std::function<void(std::string_view, Bytes)> &fn) const {
    check_poisoned();
    std::shared_lock lock(mu_);
    ensure_mapped();
    for (const auto &[key, entry] : index_) {
      if (key.starts_with(pfx)) {
        auto *ptr =
            reinterpret_cast<const std::byte *>(file_.ptr() + entry.val_offset);
        fn(key, Bytes{ptr, entry.val_len});
      }
    }
  }
 
  // --- TRANSACTION ---

  void transaction(const std::function<TxResult(Tx &)> &fn) {
    check_poisoned();
    Tx tx;
    TxResult result = fn(tx);
    if (result == rollback) {
      return;
    }
 
    std::vector<uint8_t> batch;
    batch.reserve(tx.ops_.size() * 64);
    for (const auto &op : tx.ops_) {
      detail::EntryHeader hdr{
          .flags   = op.is_delete ? detail::FLAG_TOMB : detail::FLAG_LIVE,
          .key_len = static_cast<uint8_t>(op.key.size()),
          .val_len = op.is_delete ? 0u : static_cast<uint32_t>(op.val.size()),
      };
      uint8_t raw[detail::HEADER_SIZE];
      detail::encode_header(raw, hdr);
      batch.insert(batch.end(), raw, raw + detail::HEADER_SIZE);
      batch.insert(batch.end(), op.key.begin(), op.key.end());
      if (!op.is_delete) {
        batch.insert(batch.end(), op.val.begin(), op.val.end());
      }
    }
 
    std::unique_lock lock(mu_);
 
    const size_t size_before = file_.size();
 
    if (!file_.append(batch.data(), batch.size())) {
      throw io_error("fluxen: transaction append failed");
    }
 
    if (!file_.sync()) {
      if (!file_.truncate(size_before)) {
        poisoned_ = true;
        throw poisoned_error(
            "fluxen: transaction fsync failed and truncation failed. Database "
            "file may contain a partial tail entry");
      }
      throw io_error("fluxen: transaction fsync failed");
    }
 
    size_t pos = 0;
    for (const auto &op : tx.ops_) {
      pos += detail::HEADER_SIZE + op.key.size();
      if (op.is_delete) {
        index_.erase(op.key);
      } else {
        index_[op.key] = {.val_offset = size_before + pos,
                          .val_len    = static_cast<uint32_t>(op.val.size())};
        pos += op.val.size();
      }
    }
  }
 
  // --- MAINTENANCE ---

  [[nodiscard]] bool compact() {
    check_poisoned();
    std::unique_lock lock(mu_);
 
    if (file_.is_dirty() && !file_.remap()) {
      throw io_error("fluxen: remap failed before compaction");
    }
 
    std::vector<uint8_t> buf;
    buf.reserve(file_.size());
 
    buf.insert(buf.end(), detail::MAGIC, detail::MAGIC + sizeof(detail::MAGIC));
 
    detail::IndexMap new_index;
    for (const auto &[key, entry] : index_) {
      detail::EntryHeader hdr{
          .flags   = detail::FLAG_LIVE,
          .key_len = static_cast<uint8_t>(key.size()),
          .val_len = entry.val_len,
      };
 
      uint8_t raw[detail::HEADER_SIZE];
      detail::encode_header(raw, hdr);
 
      size_t val_off = buf.size() + detail::HEADER_SIZE + key.size();
 
      buf.insert(buf.end(), raw, raw + detail::HEADER_SIZE);
      buf.insert(buf.end(), key.begin(), key.end());
 
      const auto *val_ptr = file_.ptr() + entry.val_offset;
      buf.insert(buf.end(), val_ptr, val_ptr + entry.val_len);
 
      new_index[key] = {.val_offset = val_off, .val_len = entry.val_len};
    }
 
    if (!file_.rewrite(buf)) {
      return false;
    }
 
    index_ = std::move(new_index);
    return true;
  }
 
  // --- DIAGNOSTICS ---

  [[nodiscard]] size_t key_count() const {
    check_poisoned();
    std::shared_lock lock(mu_);
    return index_.size();
  }

  [[nodiscard]] size_t file_size() const {
    check_poisoned();
    std::shared_lock lock(mu_);
    return file_.size();
  }
 
private:
  void init_file() {
    if (!file_.append(detail::MAGIC, sizeof(detail::MAGIC))) {
      throw corrupt_error("fluxen: failed to write magic header");
    }
  }

  void load_index() {
    if (file_.size() < sizeof(detail::MAGIC)) {
      throw corrupt_error("fluxen: file too small to be valid");
    }
 
    if (std::memcmp(file_.ptr(), detail::MAGIC, sizeof(detail::MAGIC)) != 0) {
      throw corrupt_error("fluxen: bad magic. File was not created by fluxen");
    }
 
    size_t pos           = sizeof(detail::MAGIC);
    size_t last_good_pos = pos;
 
    while (pos + detail::HEADER_SIZE <= file_.size()) {
      uint8_t raw[detail::HEADER_SIZE];
      std::memcpy(raw, file_.ptr() + pos, detail::HEADER_SIZE);
      detail::EntryHeader hdr = detail::decode_header(raw);
      pos += detail::HEADER_SIZE;
 
      if (pos + hdr.key_len + hdr.val_len > file_.size()) {
        break;
      }
 
      std::string key(reinterpret_cast<const char *>(file_.ptr() + pos),
                      hdr.key_len);
      pos += hdr.key_len;
 
      if (hdr.flags == detail::FLAG_TOMB) {
        index_.erase(key);
      } else {
        index_[key] = {.val_offset = pos, .val_len = hdr.val_len};
      }
 
      pos += hdr.val_len;
      last_good_pos = pos;
    }
 
    if (last_good_pos < file_.size()) {
      if (!file_.truncate(last_good_pos)) {
        throw corrupt_error(
            "fluxen: failed to truncate partial tail entry on open");
      }
    }
  }

  void append_entry(std::string_view key, const uint8_t *val, uint32_t val_len,
                    bool tombstone) {
    if (key.empty() || key.size() > detail::MAX_KEY) {
      throw key_error("fluxen: key must be between 1 and 255 bytes");
    }
 
    detail::EntryHeader hdr{
        .flags   = tombstone ? detail::FLAG_TOMB : detail::FLAG_LIVE,
        .key_len = static_cast<uint8_t>(key.size()),
        .val_len = val_len,
    };
 
    std::vector<uint8_t> buf;
    buf.resize(detail::HEADER_SIZE + key.size() + val_len);
    detail::encode_header(buf.data(), hdr);
    std::memcpy(buf.data() + detail::HEADER_SIZE, key.data(), key.size());
    if (val && val_len) {
      std::memcpy(buf.data() + detail::HEADER_SIZE + key.size(), val, val_len);
    }
 
    const size_t size_before = file_.size();
 
    if (!file_.append(buf.data(), buf.size())) {
      if (!file_.truncate(size_before)) {
        poisoned_ = true;
        throw poisoned_error(
            "fluxen: append failed and truncation failed. Database file may "
            "contain a partial tail entry");
      }
      throw io_error("fluxen: append failed");
    }
 
    if (tombstone) {
      if (auto it = index_.find(key); it != index_.end()) {
        index_.erase(it);
      }
    } else {
      index_[std::string(key)] = {.val_offset = file_.size() - val_len,
                                  .val_len    = val_len};
    }
  }

  void ensure_mapped() const {
    if (!file_.is_dirty()) {
      return;
    }
 
    std::unique_lock sync_lock(sync_mutex_);
    if (file_.is_dirty() && !file_.remap()) {
      throw io_error("fluxen: remap failed");
    }
  }
};
 
} // namespace fluxen