src/sstable/sstable_writer.cpp - master - smychkov/SStorage

src/sstable/sstable_writer.cpp
184 строки6 KB
Андрей Смычков
security: комплексные фиксы по результатам аудита
25 апр 2026, 10:00
25 апр 2026, 10:005f1be8e
#include "sstable_writer.hpp"
#include "../util/varint.hpp"

#include <cstdio>
#include <cstring>
#include <fcntl.h>
#include <unistd.h>

namespace sstorage {

//============================================================================
// Конструктор / деструктор
//============================================================================
SSTableWriter::SSTableWriter(std::string path,
                             size_t blockSize,
                             util::CompressionType compression,
                             size_t expectedKeys,
                             size_t bitsPerKey)
    : finalPath_(std::move(path)),
      tmpPath_(finalPath_ + ".tmp"),
      blockSize_(blockSize),
      compression_(compression),
      bloom_(expectedKeys, bitsPerKey) {
}

SSTableWriter::~SSTableWriter() {
    // Если не финализирован — удаляем tmp-файл
    if (opened_ && !finished_) {
        if (out_.is_open()) out_.close();
        ::unlink(tmpPath_.c_str());
    }
}

//============================================================================
// Открытие
//============================================================================
bool SSTableWriter::open() {
    out_.open(tmpPath_, std::ios::binary | std::ios::trunc);
    if (!out_) return false;
    opened_ = true;
    return true;
}

//============================================================================
// Добавление записи
//============================================================================
// Записи накапливаются в currentBlock_. Когда блок превысил blockSize_,
// он флашится на диск и начинается новый.
//============================================================================
bool SSTableWriter::add(const Record& r) {
    if (!opened_ || finished_) return false;

    if (numRecords_ == 0) {
        minKey_ = r.key();
    }
    maxKey_ = r.key();
    ++numRecords_;

    // Добавляем в Bloom filter
    bloom_.add(r.key());

    // Добавляем в текущий блок
    currentBlock_.add(r);

    // Если блок превысил таргет — флашим
    if (currentBlock_.currentSize() >= blockSize_) {
        if (!flushCurrentBlock()) return false;
    }
    return true;
}

//============================================================================
// Сброс текущего блока на диск
//============================================================================
bool SSTableWriter::flushCurrentBlock() {
    if (currentBlock_.empty()) return true;

    // 1. Финализируем raw-содержимое блока
    std::string raw = currentBlock_.finish();
    std::string lastKey = currentBlock_.lastKey();

    // 2. Сжимаем и формируем on-disk bytes
    std::string onDisk;
    buildBlockForDisk(raw, compression_, onDisk);

    // 3. Запоминаем координаты для index block
    IndexEntry entry;
    entry.lastKey = std::move(lastKey);
    entry.blockOffset = offset_;
    entry.blockSize = onDisk.size();
    indexEntries_.push_back(std::move(entry));

    // 4. Пишем на диск
    if (!writeBytes(onDisk)) return false;

    // 5. Готовим builder для следующего блока
    currentBlock_.reset();
    return true;
}

//============================================================================
// Запись bytes в файл с обновлением offset_
//============================================================================
bool SSTableWriter::writeBytes(const std::string& data) {
    out_.write(data.data(), data.size());
    if (!out_) return false;
    offset_ += data.size();
    return true;
}

//============================================================================
// Финализация
//============================================================================
// 1. Флашим последний блок (если не пустой).
// 2. Пишем index block.
// 3. Пишем Bloom filter block.
// 4. Пишем footer.
// 5. fsync() + close() + rename().
//============================================================================
bool SSTableWriter::finish() {
    if (!opened_ || finished_) return false;

    // 1. Финализируем последний блок
    if (!flushCurrentBlock()) return false;

    // 2. Index block
    uint64_t indexOffset = offset_;
    std::string indexBytes;
    util::encodeVarint(indexEntries_.size(), indexBytes);
    for (const auto& e : indexEntries_) {
        util::encodeVarint(e.lastKey.size(), indexBytes);
        indexBytes.append(e.lastKey);
        util::encodeVarint(e.blockOffset, indexBytes);
        util::encodeVarint(e.blockSize, indexBytes);
    }
    if (!writeBytes(indexBytes)) return false;
    uint64_t indexSize = indexBytes.size();

    // 3. Bloom filter
    uint64_t bloomOffset = offset_;
    std::string bloomBytes;
    bloom_.serialize(bloomBytes);
    if (!writeBytes(bloomBytes)) return false;
    uint64_t bloomSize = bloomBytes.size();

    // 4. Footer
    Footer footer;
    footer.indexOffset = indexOffset;
    footer.indexSize = indexSize;
    footer.bloomOffset = bloomOffset;
    footer.bloomSize = bloomSize;
    footer.numRecords = numRecords_;

    std::string footerBytes;
    footer.serialize(footerBytes);
    if (!writeBytes(footerBytes)) return false;

    // 5. fsync + close + rename для атомарности
    out_.flush();
    if (!out_) return false;

    // Получаем fd для fsync (из ofstream нельзя, делаем через C stdio trick)
    out_.close();

    // fsync через отдельный open — дороже, но корректно.
    // O_NOFOLLOW защищает от atomic replacement tmp-файла на symlink между
    // close(ofstream) и нашим open() здесь.
    int fd = ::open(tmpPath_.c_str(), O_RDONLY | O_NOFOLLOW);
    if (fd >= 0) {
        ::fsync(fd);
        ::close(fd);
    }

    // Атомарный rename: tmp → final
    if (::rename(tmpPath_.c_str(), finalPath_.c_str()) != 0) {
        ::unlink(tmpPath_.c_str());
        return false;
    }

    finished_ = true;
    return true;
}

}