drive.cpp

#include "drive.hpp"
#include "block_index.hpp"
#include "string_intern.hpp"
#include "block.hpp"
#include <cstdint>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/mman.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <linux/fs.h>
#include "logic_error.hpp"
#include <cstdlib>

static char* const CHAR_MAP_FAILED = reinterpret_cast<char*>(MAP_FAILED);

namespace AppenDB {

  class SpinLock {
  public:
    SpinLock(uint16_t *i) : lock(i) {
      uint16_t zero = 0, one = 1;
      while (!__atomic_compare_exchange_n(lock, &zero, one, false, __ATOMIC_ACQ_REL, __ATOMIC_ACQUIRE)) zero = 0;
    }
    ~SpinLock() {
      __atomic_store_n(lock, 0, __ATOMIC_RELEASE);
    }
  private:
    uint16_t *lock;
  };

  /*
  GLOBAL SETTINGS (128 bytes) {
    32 bits: magic number (0x278c5c3a)
    32 bits: version code (1)
    32 bits: mount state
    32 bits: page size (4k)

    32 bits: block size (128k)
    32 bits: # of blocks
    32 bits: shard for drive (zeroes for non-mask bits)
    32 bits: shard for drive mask

    32 bits: start of block entries in pages
    32 bits: start of blocks ((start of block entries in pages) + (16 * (# of blocks))) in pages
    32 bits: beginning of edit journal (128)
    32 bits: end of edit journal in pages

    80 bytes: reserved (0)
  }
  */

  struct DriveData {
    BlockIndex block_indexes[Constants::NUM_SHARD_PARTITIONS];
    Table tables[1024];
    StringInternLookupTable string_intern_table;

    int ref_count;
    int sync_fd;
    int mmap_fd;
    char* mmap_region;
    char* block_entries;
    char* blocks;
    size_t num_free_blocks;
    size_t num_blocks;
    size_t next_block_to_use;
    uint32_t shard;
    uint32_t shard_mask;
    uint16_t block_page_lock;
    bool read_only;

    size_t current_transaction_flushed_begin;
    size_t current_transaction_committed_begin;
    size_t current_transaction_begin;
    size_t current_transaction_write_pointer;

    DriveData() : ref_count(0), sync_fd(-1), mmap_fd(-1), mmap_region(CHAR_MAP_FAILED), block_entries(0), blocks(0), num_free_blocks(0), num_blocks(0), next_block_to_use(0), block_page_lock(0) {}
  } data[MAX_OPEN_DRIVES];

  Drive Drive::get_any_drive(){
    Drive rv;

    for(size_t i = 0; i != MAX_OPEN_DRIVES; ++i){
      if( data[i].ref_count ){
        rv._ = data + i;
        rv._add_ref();
        return rv;
      }
    }
    emit_logic_error("No drive found");
  }

  static DriveData* find_free_data_slot(){
    for(size_t i = 0; i != MAX_OPEN_DRIVES; ++i){
      if( __sync_bool_compare_and_swap (&(data[i].ref_count), 0, 1) ) return data + i;
    }
    return 0;
  }

  Block get_block(const char* blocks, uint32_t index){
    return Block(blocks + (16 * index));
  }

  char* Drive::drive_journal_start() const { return disk() + DRIVE_JOURNAL_START; }
  char* Drive::drive_journal_end() const { assert(_); return reinterpret_cast<DriveData *>(_)->block_entries; }
  BlockIndex &Drive::block_index(size_t i) { assert(_); return reinterpret_cast<DriveData *>(_)->block_indexes[i]; }
  const BlockIndex &Drive::block_index(size_t i) const { assert(_); return reinterpret_cast<DriveData *>(_)->block_indexes[i]; }
  StringInternLookupTable &Drive::string_intern_lookup() { assert(_); return reinterpret_cast<DriveData *>(_)->string_intern_table; }
  const StringInternLookupTable &Drive::string_intern_lookup() const { assert(_); return reinterpret_cast<DriveData *>(_)->string_intern_table; }
  char* Drive::disk() const { assert(_); return reinterpret_cast<DriveData *>(_)->mmap_region; }
  char* Drive::block_entries() const { assert(_); return reinterpret_cast<DriveData *>(_)->block_entries; }
  char* Drive::blocks() const { assert(_); return reinterpret_cast<DriveData *>(_)->blocks; }
  size_t& Drive::current_transaction_flushed_begin(){ assert(_); return reinterpret_cast<DriveData *>(_)->current_transaction_flushed_begin; }
  size_t& Drive::current_transaction_committed_begin(){ assert(_); return reinterpret_cast<DriveData *>(_)->current_transaction_committed_begin; }
  size_t& Drive::current_transaction_begin() { assert(_); return reinterpret_cast<DriveData *>(_)->current_transaction_begin; }
  size_t& Drive::current_transaction_write_pointer() { assert(_); return reinterpret_cast<DriveData *>(_)->current_transaction_write_pointer; }
  Block Drive::get_block(uint32_t block_index) const { assert(_); return Block(reinterpret_cast<DriveData *>(_)->block_entries + (16 * block_index)); }

  void Drive::add_table(Table table){
    // TODO: update block_index so that it initializes the table header.
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    TableIndex table_index = table.index();
    assert(table_index < 1023);
    assert(table_index >= 4);
    size_t const new_table_size = table.representation_size();

    // TODO: fail the transaction if the table exists and doesn't match exactly
    assert(!_data.tables[table_index]);

    uint32_t block_index;
    const char* table_location_on_disk;

    // scan over all blocks currently used for tables
    for(size_t i = 0;; ++i){

      // allocate a new block because no existing one has room
      if(i == 1024){
        block_index = find_free_block(new_table_size);
        table_location_on_disk = block_data_end(block_index);
        block_allocate(block_index, 0x3FF, 0, table.representation_begin(), new_table_size, 0, 0);
        break;
      }

      // check to see if this table's block has room
      Table const t = _data.tables[table_index];
      if(!t) continue;
      const char * const table_begin = t.representation_begin();
      size_t const table_offset = table_begin - _data.blocks;
      uint32_t table_block_index = table_offset / DRIVE_BLOCK_SIZE;

      if(block_has_room(table_block_index, new_table_size)){
        block_index = table_block_index;
        table_location_on_disk = block_data_end(block_index);
        block_append_data(block_index, table.representation_begin(), new_table_size, 0, 0);
        break;
      }
    }

    // add the table to _data.tables
    _data.tables[table_index] = Table(table_location_on_disk);
  }

  void Drive::allocate_block_page(uint32_t block_container[256], uint32_t &num_blocks) {
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    SpinLock lock(&_data.block_page_lock);
    num_blocks = 0;
    uint32_t end = ::std::min(_data.num_blocks, ((_data.next_block_to_use / Constants::BLOCKS_PER_PAGE) + 1) * Constants::BLOCKS_PER_PAGE);
    for (uint32_t i = _data.next_block_to_use; i < end; ++i) {
      Block block = get_block(i);
      if(block.is_free()){
        block_container[num_blocks] = i;
        num_blocks++;
      }
    }
    _data.next_block_to_use = end;
  }

  void Drive::block_free(uint32_t block_index){
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    //Block block = get_block(block_index);
    // TODO: This should only be called from a shard thread
    //this->block_index().remove_block(block.type(), block.shard_filter(), block_index);

    for(;;){
      // TODO: do this as one larger write instead of lots of smaller writes
      write_data(reinterpret_cast<const char *>(&FREE_BLOCK), 16, (_data.block_entries - _data.mmap_region) + (16 * block_index));
      ++block_index;
      if(!get_block(block_index).is_continuation()) break;
    }
  }

  void Drive::block_allocate(uint32_t block_index, TableIndex type, uint16_t shard_mask, const char *data, size_t data_length, uint32_t min_time, uint32_t max_time){
    char* target_pointer;
    size_t target_pointer_capacity;
    block_allocate_pointer(block_index, type, shard_mask, data_length, min_time, max_time, target_pointer, target_pointer_capacity);

    // TODO: handle cases where we have to do writes to multiple locations
    assert(data_length == target_pointer_capacity);

    ::memcpy(target_pointer, data, data_length);
  }

  // TODO: Add a way to allocate a 4k block of blocks.
  void Drive::block_allocate_pointer(uint32_t block_index, TableIndex type, uint16_t shard_mask, size_t data_length, uint32_t min_time, uint32_t max_time, char*& target_pointer, size_t& target_pointer_capacity){
    uint32_t const base_block_index = block_index;
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    assert(get_block(block_index).is_free());
    assert(data_length);
    const size_t num_blocks_needed = (data_length + DRIVE_BLOCK_SIZE - 1) / DRIVE_BLOCK_SIZE;

    // write to the block entries
    uint64_t temp[2 * num_blocks_needed];
    temp[0] = (static_cast<uint64_t>(type << 6)) | (static_cast<uint64_t>(shard_mask) << 16) | (static_cast<uint64_t>(data_length) << 32);
    temp[1] = (static_cast<uint64_t>(min_time)) | (static_cast<uint64_t>(max_time) << 32);
    for(size_t i = 1; i != num_blocks_needed; ++i){
      assert(get_block(block_index + i).is_free());
      temp[2*i] = 1 | ((i) << 1) | (static_cast<uint64_t>(0xFFFFFFFF) << 32);
      temp[2*i + 1] = (static_cast<uint64_t>(0xFFFFFFFF) << 32);
    }
    write_data(reinterpret_cast<char *>(temp), 16 * num_blocks_needed, (_data.block_entries - _data.mmap_region) + (16 * block_index));
    assert(!get_block(block_index).is_free());

    // write the data to the blocks
    write_data_pointer(data_length, (_data.blocks - _data.mmap_region) + (DRIVE_BLOCK_SIZE * base_block_index), target_pointer, target_pointer_capacity);

  }

  void Drive::block_allocate_sampled(uint32_t block_index, TableIndex type, uint16_t shard_mask, uint16_t sampled_shard_mask, const char *data, size_t data_length, uint32_t min_time, uint32_t max_time){
    char* target_pointer;
    size_t target_pointer_capacity;
    block_allocate_sampled_pointer(block_index, type, shard_mask, sampled_shard_mask, data_length, min_time, max_time, target_pointer, target_pointer_capacity);

    // TODO: handle cases where we have to do writes to multiple locations
    assert(data_length == target_pointer_capacity);

    ::memcpy(target_pointer, data, data_length);
  }

  void Drive::block_allocate_sampled_pointer(uint32_t block_index, TableIndex type, uint16_t shard_mask, uint16_t sampled_shard_mask, size_t data_length, uint32_t min_time, uint32_t max_time, char*& target_pointer, size_t& target_pointer_capacity){
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    assert(get_block(block_index).is_free());
    assert(get_block(block_index + 1).is_free());
    const size_t num_blocks_needed = (data_length <= (2 * DRIVE_BLOCK_SIZE)) ? 2 : ((data_length + DRIVE_BLOCK_SIZE - 1) / DRIVE_BLOCK_SIZE);

    // write to the block entries
    uint64_t temp[2*num_blocks_needed];
    temp[0] = (static_cast<uint64_t>(type << 6)) | (static_cast<uint64_t>(shard_mask) << 16) | (static_cast<uint64_t>(data_length) << 32);
    temp[1] = (static_cast<uint64_t>(min_time)) | (static_cast<uint64_t>(max_time) << 32);
    temp[2] = 3 | (static_cast<uint64_t>(sampled_shard_mask) << 16) | (static_cast<uint64_t>(0xFFFFFFFF) << 32);
    temp[3] = (static_cast<uint64_t>(0xFFFFFFFF) << 32);
    for(size_t i = 2; i != num_blocks_needed; ++i){
      assert(get_block(block_index + i).is_free());
      temp[i*2] = 1 | ((i) << 1) | (static_cast<uint64_t>(0xFFFFFFFF) << 32);
      temp[i*2 + 1] = (static_cast<uint64_t>(0xFFFFFFFF) << 32);
    }
    write_data(reinterpret_cast<char *>(temp), 16 * num_blocks_needed, (_data.block_entries - _data.mmap_region) + (16 * block_index));
    assert(!get_block(block_index).is_free());
    assert(!get_block(block_index + 1).is_free());

    // write the data to the blocks
    write_data_pointer(data_length, (_data.blocks - _data.mmap_region) + (DRIVE_BLOCK_SIZE * block_index), target_pointer, target_pointer_capacity);

  }

  void Drive::block_append_data(uint32_t block_index, const char *data, size_t data_length, uint32_t new_min_time, uint32_t new_max_time){
    char* target_pointer;
    size_t target_pointer_capacity;
    block_append_data_pointer(block_index, data_length, new_min_time, new_max_time, target_pointer, target_pointer_capacity);

    // TODO: handle cases where we have to do writes to multiple locations
    assert(data_length == target_pointer_capacity);

    ::memcpy(target_pointer, data, data_length);
  }

  void Drive::block_append_data_pointer(uint32_t block_index, size_t data_length, uint32_t new_min_time, uint32_t new_max_time, char*& target_pointer, size_t& target_pointer_capacity){
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    assert(!get_block(block_index).is_free());
    size_t const current_size = get_block(block_index).data_length();
    size_t const total_data_size = data_length + current_size;

    const uint32_t block_index_end = 1 + block_index + ((total_data_size - 1) / DRIVE_BLOCK_SIZE);
    const uint32_t base_block_index = block_index;

    // update the data_length and max_time for the root block entry
    {
      char temp[16];
      ::memcpy(temp, _data.block_entries + (16 * block_index), 16);
      *reinterpret_cast<uint32_t *>(temp + 4) = total_data_size;
      *reinterpret_cast<uint32_t *>(temp + 8) = new_min_time;
      *reinterpret_cast<uint32_t *>(temp + 12) = new_max_time;
      write_data(temp, 16, (_data.block_entries - _data.mmap_region) + (16 * block_index));
    }

    // mark any additional non-root blocks as part of the block
    for(;;){
      ++block_index;
      if(block_index == block_index_end) break;
      Block block(_data.block_entries + (16 * block_index));
      if(!block.is_continuation()){
        const size_t num_new_blocks_to_write = block_index_end - block_index;
        char temp[16 * num_new_blocks_to_write];
        for(uint32_t i = 0; i != num_new_blocks_to_write; ++i){
          assert((Block(_data.block_entries + (16 * (block_index + i)))).is_free());
          *reinterpret_cast<uint64_t *>(temp + (i*16)) = 1 | ((block_index + i - base_block_index) << 1) | (static_cast<uint64_t>(0xFFFFFFFF) << 32);
          *reinterpret_cast<uint64_t *>(temp + (i*16) + 8) = (static_cast<uint64_t>(0xFFFFFFFF) << 32);
        }
        write_data(temp, 16 * num_new_blocks_to_write, (_data.block_entries - _data.mmap_region) + (16 * block_index));
        break;
      }
    }

    // write the data into the block
    write_data_pointer(data_length, (_data.blocks - _data.mmap_region) + (DRIVE_BLOCK_SIZE * base_block_index) + current_size, target_pointer, target_pointer_capacity);
  }

  void Drive::block_truncate_data(uint32_t block_index, size_t new_data_length){
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    assert(!get_block(block_index).is_free());

    // update the data_length and max_time for the root block entry
    {
      char temp[16];
      ::memcpy(temp, _data.block_entries + (16 * block_index), 16);
      *reinterpret_cast<uint32_t *>(temp + 4) = new_data_length;
      write_data(temp, 16, (_data.block_entries - _data.mmap_region) + (16 * block_index));
    }
  }

  const char* Drive::block_data_end(uint32_t block_index) const {
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    size_t const current_size = get_block(block_index).data_length();
    return _data.blocks + (DRIVE_BLOCK_SIZE * block_index) + current_size;
  }

  const char* Drive::block_data_start(uint32_t block_index) const {
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    return _data.blocks + (DRIVE_BLOCK_SIZE * block_index);
  }

  uint32_t Drive::find_free_block(size_t capacity) const {
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    SpinLock lock(&_data.block_page_lock);
    const size_t num_blocks_needed = (capacity <= (DRIVE_BLOCK_SIZE)) ? 1 : ((capacity + DRIVE_BLOCK_SIZE - 1) / DRIVE_BLOCK_SIZE);

    size_t next_block_to_use = _data.next_block_to_use;
    const size_t num_blocks = _data.num_blocks;
    const size_t initial_next_block_to_use = next_block_to_use;
    size_t num_consecutive = 0;
    for(;;){
      char* block_entry_location = _data.block_entries + (16 * next_block_to_use);
      Block block(block_entry_location);
      if(block.is_used()){
        ++next_block_to_use;
        if(next_block_to_use == num_blocks) next_block_to_use = 0;
        if(next_block_to_use == initial_next_block_to_use){
          // FIXME: we're out of space
          fatal_error("We're out of space on our drive");
        }
      }

      // we found a free block, check to see if it is long enough
      // if it is, we're good to go
      else if(++num_consecutive == num_blocks_needed) {
        _data.next_block_to_use = next_block_to_use;
        return next_block_to_use + 1 - num_blocks_needed;
      }
    }
  }

  uint32_t Drive::find_free_block() const {
    return find_free_block(DRIVE_BLOCK_SIZE);
  }

  size_t Drive::block_data_size(uint32_t block_index) const {
    assert(_);
    Block block = get_block(block_index);
    return block.data_length();
  }

  size_t Drive::block_capacity(uint32_t block_index, size_t max_size_to_scan) const {
    if(DRIVE_BLOCK_SIZE_MAX < max_size_to_scan) max_size_to_scan = DRIVE_BLOCK_SIZE_MAX;
    assert(_);
    size_t total_size = DRIVE_BLOCK_SIZE;
    ++block_index;
    for(;;){
      if(total_size >= max_size_to_scan) return total_size;

      // not necessary because we have a termination block at the end
      // if(block_index == num_blocks()) return total_size;

      Block block = get_block(block_index);
      if(block.is_continuation() || block.is_free()){
        ++block_index;
        total_size += DRIVE_BLOCK_SIZE;
      }

      // no more free space
      else return total_size;
    }
  }

  size_t Drive::block_space_available(uint32_t block_index, size_t max_size_of_additional_data_to_scan) const {
    size_t const block_current_used = block_data_size(block_index);
    size_t const needed_new_size = max_size_of_additional_data_to_scan + block_current_used;
    return block_capacity(block_index, needed_new_size) - block_current_used;
  }

  bool Drive::block_has_room(uint32_t block_index, size_t size_of_additional_data) const {
    size_t const needed_new_size = size_of_additional_data + block_data_size(block_index);
    return block_capacity(block_index, needed_new_size) >= needed_new_size;
  }

  void Drive::attach(const char *file_name, bool override_mounted, bool read_only){
    _remove_ref();

    DriveData* d = find_free_data_slot();
    if(!d) return;
    DriveData &_data = *d;
    _data.read_only = read_only;
    _ = &_data;
    for (size_t i = 0; i < Constants::NUM_SHARD_PARTITIONS; ++i) {
      _data.block_indexes[i].drive = *this;
    }

    if(Drive::attach(
      file_name,
      _data.block_indexes,
      _data.string_intern_table,
      _data.tables,
      _data.sync_fd,
      _data.mmap_fd,
      _data.mmap_region,
      _data.block_entries,
      _data.blocks,
      _data.num_free_blocks,
      _data.num_blocks,
      _data.next_block_to_use,
      _data.shard,
      _data.shard_mask,
      override_mounted,
      _data.read_only
    )){
      return;
    }
    else {
      _ = 0;
      _data.ref_count = 0;
      return;
    }
  }

  void Drive::detach(){
    _remove_ref();
  }

  size_t Drive::num_blocks() const {
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    return _data.num_blocks;
  }

  size_t Drive::num_free_blocks() const {
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    return _data.num_free_blocks;
  }

  uint32_t Drive::shard_mask() const {
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    return _data.shard_mask;
  }

  uint32_t Drive::shard() const {
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    return _data.shard;
  }

  Table Drive::get_table(TableIndex table_index) const {
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    assert(_data.tables[table_index]);
    return _data.tables[table_index];
  }

  InputTable Drive::get_input_table(TableIndex table_index) const {
    assert(_);
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    assert(_data.tables[table_index]);
    return InputTable(_data.tables[table_index].representation_begin());
  }

  // NOTE: _add_ref can assume that the reference to the entry is at least 1 during the duration of its execution due to the existence of a valid Drive object with the same underlying data
  void Drive::_add_ref(){
    if(!_) return;
    DriveData &_data = *reinterpret_cast<DriveData *>(_);
    __atomic_add_fetch(&_data.ref_count, 1, __ATOMIC_RELEASE);
  }

  void Drive::_remove_ref(){
    if(!_) return;
    DriveData &_data = *reinterpret_cast<DriveData *>(_);

    int ref_count = __atomic_sub_fetch(&_data.ref_count, 1, __ATOMIC_RELEASE);
    if (ref_count == 0) {
      Drive::detach(
        _data.block_indexes,
        _data.string_intern_table,
        _data.tables,
        _data.sync_fd,
        _data.mmap_fd,
        _data.mmap_region,
        _data.block_entries,
        _data.blocks,
        _data.num_free_blocks,
        _data.num_blocks,
        _data.next_block_to_use,
        _data.shard,
        _data.shard_mask,
        _data.read_only
      );
    }
    _ = 0;
  }
}