format.cpp

#include "drive.hpp"
#include "block_index.hpp"
#include "string_intern.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 "constants.hpp"
#include <cstdlib>
#include <iostream>

static uint64_t get_file_size(int fd){
  uint64_t file_size = 0;
  int ioctl_rv = ioctl(fd, BLKGETSIZE64, &file_size);
  assert(!ioctl_rv);
  assert(file_size);
  return file_size;
}

static const size_t BLOCK_ENTRY_SIZE = 16;
static const uint32_t DRIVE_MOUNTED_STATE = 1;
static const uint64_t STRING_INTERN_COUNT_MEMORY_LENGTH = (AppenDB::Constants::DEFAULT_NUM_SHARDS << 1);

static const char KEY_VALUE_TABLE_DEFINITION[] = {
  0x40, 0x00, 0x00, 0x00, 0x00, 0x03, 0x18, 0x00, // table header:
                                                  //   table index: 1
                                                  //   num fields:  2
                                                  //   table name offset: 0x18
  0x48, 0x04, 0x13, 0x00, // first field: "time"
                          //   special time type
                          //
  0x0C, 0x00, 0x13, 0x00, // second field: "key"
                          //   not nullable or special
                          //   indexed and unique
  0x01, 0x00, 0x12, 0x00, // third field: "value"
                          //   nullable
                          //   not special, indexed or unique
  0x40, 0x00, 0x13, 0x00, // field termination entry
   'K',  '/',  'V',  't',
   'i',  'm',  'e',  'k',
   'e',  'y',  'v',  'a',
   'l',  'u',  'e'
};

static const char STRING_INTERNMENT_TABLE_DEFINITION[] = {
  0x80, 0x00, 0x00, 0x00, 0x00, 0x02, 0x14, 0x00,
  0x48, 0x04, 0x1c, 0x00,
  0x0c, 0x00, 0x1c, 0x00,
  0x40, 0x00, 0x1b, 0x00,
  'S', 't', 'r', 'i', 'n', 'g', 'I', 'n', 't', 'e', 'r', 'n', 'm', 'e', 'n', 't',
  't', 'i', 'm', 'e',
  'k', 'e', 'y'
};


static const char ENGAGEMENTS_TABLE_DEFINITION[] = {
  0x00, 0x01, 0x00, 0x00, 0x00, 0x14, 0x5c, 0x00,
  0x48, 0x04, 0x5f, 0x00,
  0x04, 0x08, 0x5f, 0x00,
  0x05, 0x08, 0x65, 0x00,
  0x05, 0x08, 0x6c, 0x00,
  0x01, 0x08, 0x6f, 0x00,
  0x01, 0x08, 0x71, 0x00,
  0x01, 0x08, 0x74, 0x00,
  0x01, 0x08, 0x7b, 0x00,
  0x01, 0x08, 0x85, 0x00,
  0x01, 0x04, 0x8b, 0x00,
  0x01, 0x10, 0x93, 0x00,
  0x01, 0x02, 0x9b, 0x00,
  0x21, 0x00, 0x9e, 0x00,
  0x01, 0x02, 0xa4, 0x00,
  0x01, 0x02, 0xa9, 0x00,
  0x00, 0x02, 0xab, 0x00,
  0x01, 0x04, 0xae, 0x00,
  0x01, 0x04, 0xaf, 0x00,
  0x01, 0x04, 0xad, 0x00,
  0x01, 0x04, 0xb5, 0x00,
  0x40, 0x00, 0xb6, 0x00,
'e', 'n', 'g', 'a', 'g', 'e', 'm', 'e', 'n', 't', 's',
't', 'i', 'm', 'e',
'a', 'c', 'c', 'o', 'u', 'n', 't', '_', 'i', 'd',
'c', 'r', 'e', 'a', 't', 'i', 'v', 'e', '_', 'i', 'd',
'a', 'd', '_', 's', 'l', 'o', 't',
'a', 'd', '_', 'b', 'i', 'd',
'p', 'a', 'g', 'e', '_', 'i', 'd',
'r', 'e', 'f', 'e', 'r', 'r', 'e', 'r', '_', 'i', 'd',
'i', 'n', 't', 'e', 'r', 'a', 'c', 't', 'i', 'o', 'n', '_', 'i', 'd',
'r', 'e', 'q', 'u', 'e', 's', 't', '_', 'i', 'd',
'i', 'p', 'v', '4', '_', 'a', 'd', 'd', 'r', 'e', 's', 's',
'i', 'p', 'v', '6', '_', 'a', 'd', 'd', 'r', 'e', 's', 's',
'c', 'o', 'u', 'n', 't', 'r', 'y',
'u', 's', 'e', 'r', '_', 'a', 'g', 'e', 'n', 't',
'p', 'l', 'a', 'c', 'e', 'm', 'e', 'n', 't',
'd', 'e', 'v', 'i', 'c', 'e',
's', 'u', 'b', 't', 'y', 'p', 'e',
's', 'p', 'e', 'n', 't',
'x', 'y',
'w', 'i', 'd', 't', 'h', '_', 'h', 'e', 'i', 'g', 'h', 't',
'f', 'l', 'a', 'g', 's'
};

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

  32 bits: block size (256k)
  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 (8k) in pages
  32 bits: start of blocks (8k + (16 * (# of blocks))) in pages
  32 bits: beginning of edit journal (128)
  32 bits: end of edit journal (8k) in pages

  80 bytes: reserved (0)
}
*/

namespace AppenDB {
  bool Drive::format(const char *file_name){
    int fd = -1;
    unsigned int num_errors = 0;

    for(;;){
      fd = open(file_name, O_CLOEXEC | O_NOATIME | O_RDWR);
      if(fd == -1){
        if(errno == EINTR) continue;
        else if(num_errors > 2) return false;
        else ++num_errors;
      }
      else break;
    }

    uint64_t file_size = get_file_size(fd);
    if(file_size & (DRIVE_PAGE_SIZE - 1)) return false;
    char* mmap_region = reinterpret_cast<char*>(mmap(0, file_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0));

    const uint32_t block_entries_offset_in_pages = 1; // allocate the rest of the first page for an index of WAL blocks
    const uint32_t num_blocks = (file_size - DRIVE_PAGE_SIZE * block_entries_offset_in_pages) / (DRIVE_BLOCK_SIZE + 16);

    // blocks will start after:
    //   the block entry edits
    //   room for a 16 byte block entry & a termination entry rounded up to the next page
    const uint32_t blocks_offset_in_pages = block_entries_offset_in_pages + 1 + (((num_blocks + 1) * 16) - 1) / DRIVE_PAGE_SIZE;

    // write header and all the block entry edits
    uint32_t* __restrict mmap_region_32 = reinterpret_cast<uint32_t*>(mmap_region);
    mmap_region_32[0] = DRIVE_MAGIC_NUMBER;
    mmap_region_32[1] = DRIVE_VERSION;
    mmap_region_32[2] = 2;
    mmap_region_32[3] = DRIVE_PAGE_SIZE;
    mmap_region_32[4] = DRIVE_BLOCK_SIZE;
    mmap_region_32[5] = num_blocks;
    mmap_region_32[6] = 0;
    mmap_region_32[7] = 0;
    mmap_region_32[8] = block_entries_offset_in_pages;
    mmap_region_32[9] = blocks_offset_in_pages;
    mmap_region_32[10] = DRIVE_JOURNAL_START;
    mmap_region_32[11] = block_entries_offset_in_pages;
    ::memset(mmap_region_32 + 12, 0, block_entries_offset_in_pages << 12);

    // write the block entries
    __uint128_t* __restrict block_entries = reinterpret_cast<__uint128_t*>(mmap_region + (block_entries_offset_in_pages * DRIVE_PAGE_SIZE));
    block_entries[0] = 0x0002FFC0ULL; // special first block should be unused and marked as special/reserved
    block_entries[1] = 0x0000FFC0ULL + (sizeof(KEY_VALUE_TABLE_DEFINITION) << 32); // allocate space for table definitions & set the initial size for the K/V table
    block_entries[2] = 0x0000FFC0ULL + (sizeof(STRING_INTERNMENT_TABLE_DEFINITION) << 32); // allocate space for table definitions & set the initial size for the StringIntern table
    ::std::fill(block_entries + 3, block_entries + num_blocks, FREE_BLOCK);
    block_entries[4] = 0x0000FFC0ULL + (sizeof(ENGAGEMENTS_TABLE_DEFINITION) << 32); // allocate space for table definitions & set the initial size for the table
    block_entries[num_blocks] = TERMINATION_BLOCK;

    // write the actual data for the K/V table
    ::memcpy(mmap_region + (blocks_offset_in_pages * DRIVE_PAGE_SIZE) + (1 * DRIVE_BLOCK_SIZE), KEY_VALUE_TABLE_DEFINITION, sizeof(KEY_VALUE_TABLE_DEFINITION));
    ::memcpy(mmap_region + (blocks_offset_in_pages * DRIVE_PAGE_SIZE) + (2 * DRIVE_BLOCK_SIZE), STRING_INTERNMENT_TABLE_DEFINITION, sizeof(STRING_INTERNMENT_TABLE_DEFINITION));
    ::memcpy(mmap_region + (blocks_offset_in_pages * DRIVE_PAGE_SIZE) + (4 * DRIVE_BLOCK_SIZE), ENGAGEMENTS_TABLE_DEFINITION, sizeof(ENGAGEMENTS_TABLE_DEFINITION));

    // unmount
    if(
      fdatasync(fd) && fdatasync(fd) // write all data with the mount guard active (this ensures that everything is written)
    ) fatal_error("Error syncing the data prior to marking the drive as unmounted");
    mmap_region_32[2] = 0; // Now that we know all the data is synced, it is safe to unmount
    if(
      fdatasync(fd) && fdatasync(fd) // write out that the drive has been unmounted since it is now ready to be read
    ) fatal_error("Error syncing the data after marking the drive as unmounted");
    if(munmap(mmap_region, file_size)) fatal_error("Error unmapping the drive");
    if(
      close(fd) && close(fd) && close(fd) && close(fd) && close(fd) // try to close the file 5 times
    ) fatal_error("Error closing the file descriptor for the drive");

    return true;
  }

  bool Drive::attach(
    const char *file_name,
    BlockIndex block_indexes[Constants::NUM_SHARD_PARTITIONS],
    StringInternLookupTable &string_intern_lookup,
    Table *tables,
    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,
    bool override_mounted,
    const bool &read_only
  ){
    uint32_t block_entries_offset_in_pages;
    uint32_t blocks_offset_in_pages;
    uint64_t file_size;
    unsigned int num_errors;
    uint32_t num_needed_block_entry_pages;
    uint32_t num_needed_block_pages;
    const uint32_t* global_buffer;

    num_errors = 0;
    for(;;){
      sync_fd = open(file_name, O_CLOEXEC | O_NOATIME | (read_only ? O_RDONLY : O_RDWR) | O_DSYNC);
      if(sync_fd == -1){
        if(errno == EINTR) continue;
        else if(num_errors > 2) goto on_error;
        else ++num_errors;
      }
      else break;
    }

    file_size = get_file_size(sync_fd);

    num_errors = 0;
    for(;;){
      mmap_fd = open(file_name, O_CLOEXEC | O_NOATIME | (read_only ? O_RDONLY : O_RDWR));
      if(mmap_fd == -1){
        if(errno == EINTR) continue;
        else if(num_errors > 2) goto on_error;
        else ++num_errors;
      }
      else break;
    }

    num_errors = 0;
    for(;;){
      mmap_region = reinterpret_cast<char*>(mmap(0, file_size, PROT_READ | (read_only ? 0 : PROT_WRITE), MAP_SHARED, mmap_fd, 0));
      if(mmap_region == MAP_FAILED){
        if(errno == EINTR) continue;
        else if(num_errors > 2) goto on_error;
        else ++num_errors;
      }
      else break;
    }

    global_buffer = reinterpret_cast<const uint32_t*>(mmap_region);

    if(
      (global_buffer[0] != DRIVE_MAGIC_NUMBER) ||
      (global_buffer[1] != DRIVE_VERSION) ||
      ((global_buffer[2] != 0) && (!override_mounted)) ||
      (global_buffer[3] != DRIVE_PAGE_SIZE) ||
      (global_buffer[4] != DRIVE_BLOCK_SIZE)
    ) goto on_error;

    num_blocks = global_buffer[5];
    shard = global_buffer[6];
    shard_mask = global_buffer[7];
    block_entries_offset_in_pages = global_buffer[8];
    blocks_offset_in_pages = global_buffer[9];
    num_needed_block_entry_pages = ((num_blocks * BLOCK_ENTRY_SIZE) + (DRIVE_PAGE_SIZE - 1)) / DRIVE_PAGE_SIZE;
    num_needed_block_pages = (num_blocks * DRIVE_BLOCK_SIZE) / DRIVE_PAGE_SIZE;

    if(
      (blocks_offset_in_pages < (block_entries_offset_in_pages + num_needed_block_entry_pages)) ||
      (((blocks_offset_in_pages + num_needed_block_pages) * DRIVE_PAGE_SIZE) > (file_size)) ||
      (global_buffer[10] != DRIVE_JOURNAL_START) ||
      (global_buffer[11] != block_entries_offset_in_pages)
    ) goto on_error;

    for(size_t i = 12; i != 32; ++i){
      if(global_buffer[i] != 0) goto on_error;
    }

    // let the kernel know that we're going to be reading in the block entries
    {
      const size_t readahead_size = ((BLOCK_ENTRY_SIZE * num_blocks) + DRIVE_PAGE_SIZE) % DRIVE_PAGE_SIZE;
      madvise(mmap_region + ( block_entries_offset_in_pages * DRIVE_PAGE_SIZE), readahead_size, MADV_WILLNEED);
    }

    // set up some pointers
    block_entries = mmap_region + (block_entries_offset_in_pages * DRIVE_PAGE_SIZE);
    blocks = mmap_region + (blocks_offset_in_pages * DRIVE_PAGE_SIZE);

    // ensure that we don't have multiple clients mount the drive
    // write that we've mounted this drive
    // NOTE: This extra check has a race condition. There is no way around this,
    // but the check is redundant since no drive should ever be mounted multiple times anyway
    if(!read_only){
      if(4 != ::pwrite(sync_fd, &DRIVE_MOUNTED_STATE, 4, 8)) goto on_error;
    }

    // apply journal edits
    if(!read_only) Drive::apply_journal(mmap_region, mmap_region + DRIVE_JOURNAL_START, block_entries);

    // zero out all tables
    for(size_t i = 0; i != 1024; ++i) tables[i] = 0;

    // scan over all blocks
    {
      // HEADER is 2^10 possible tables * (2^16 possible shards - 2^x number of shard partitions)
      const size_t block_index_length = ((1 << (26 - (Constants::NUM_SHARD_BITS - 1))) + (num_blocks << 1) + 0x3FF) & 0xFFFFFC00; // round up to the next page (this is in uint32_t size)

      void* block_index_memory = mmap(0, Constants::NUM_SHARD_PARTITIONS * block_index_length * sizeof(uint32_t), PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0);
      BlockIndexImport block_import;
      assert(block_index_memory != MAP_FAILED);
      uint32_t *string_intern_counts_memory = reinterpret_cast<uint32_t *>(mmap(0, STRING_INTERN_COUNT_MEMORY_LENGTH * sizeof(uint32_t), PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0));
      string_intern_lookup.initialize(string_intern_counts_memory, string_intern_counts_memory + (STRING_INTERN_COUNT_MEMORY_LENGTH >> 1));

      const uint64_t* __restrict block_entries_64 = reinterpret_cast<uint64_t *>(block_entries);

      // check that the beginning of the drive has a special reserved block then a table block
      assert(0xFFC0 == (((reinterpret_cast<uint64_t *>(block_entries))[0 * 2]) & 0xFFC0));
      assert(2 == ((((reinterpret_cast<uint64_t *>(block_entries))[0 * 2]) >> 16) & 0xFFFF));
      assert(0xFFC0 == (((reinterpret_cast<uint64_t *>(block_entries))[1 * 2]) & 0xFFC0));
      assert(0 == ((((reinterpret_cast<uint64_t *>(block_entries))[1 * 2]) >> 16) & 0xFFFF));

      for(size_t i = 0; i != num_blocks; ++i){
        const uint64_t block_entry_low = block_entries_64[i * 2];
        if(block_entry_low & 1) continue; // skip continuation blocks
        const size_t data_length = block_entry_low >> 32;
        const unsigned int type = block_entry_low & 0xFFC0;

        // special block found
        if(type == 0xFFC0){
          const unsigned int code = (block_entry_low >> 16) & 0xFFFF;

          // free
          if(code == 1){
            ++num_free_blocks;
            block_import.add_free_block(i);
          }

          // table metadata
          else if(code == 0){
            const char* __restrict tables_source = blocks + (DRIVE_BLOCK_SIZE * i);
            const char* __restrict tables_source_end = tables_source + data_length;
            while(tables_source != tables_source_end){
              assert(tables_source < tables_source_end);
              const TableIndex table_index = (*reinterpret_cast<const uint32_t*>(tables_source) >> 6) & 0x3FF;
              assert(!tables[table_index]);
              tables[table_index] = tables_source;
              // TODO: consider just passing Table *tables into block_index
              block_import.tables_used[table_index] = true;
              assert(tables[table_index]);
              tables_source += tables[table_index].representation_size();
            }

            // TODO: add this block in the index of table blocks
            // TODO: add data into table
          }

          // other special block
          else {
            assert(code == 2);
          }
        }

        // normal, used block found
        else {
          if (((block_entry_low >> 6) & 0x3FF) == Constants::STRING_INTERNMENT_TABLE_INDEX) {
            string_intern_lookup.add_block(i);
          }
          block_import.add_block(i, block_entry_low, block_entries[i * 2 + 1]);
        }
      }
      next_block_to_use = block_import.next_block_to_use();
      for (size_t i = 0; i < Constants::NUM_SHARD_PARTITIONS; ++i) {
        block_indexes[i].initialize(
          reinterpret_cast<uint32_t*>(block_index_memory) + (i * block_index_length),
          block_index_length,
          i,
          block_import
        );
      }

    }
    return true;

  on_error:
    if((sync_fd != -1) && (close(sync_fd) != 0) && (close(sync_fd) != 0)) goto fatal;
    if((mmap_fd != -1) && (close(mmap_fd) != 0) && (close(mmap_fd) != 0)) goto fatal;
    if((mmap_region != MAP_FAILED) && (munmap(mmap_region, file_size)) && (munmap(mmap_region, file_size))) goto fatal;
    return false;

  fatal:
    // TODO: possibly handle this better
    fatal_error("Fatal Error while trying to mount a drive");
  }

  void Drive::detach(
    BlockIndex block_indexes[Constants::NUM_SHARD_PARTITIONS],
    StringInternLookupTable &string_intern_lookup,
    Table *tables,
    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,
    const bool &read_only
  ){
    if(mmap_fd != -1){
      uint64_t file_size = get_file_size(mmap_fd);

      // make sure that all data is flushed
      if(!read_only && fdatasync(mmap_fd) && fdatasync(mmap_fd) && fdatasync(mmap_fd)) goto fatal;

      // clear out the journal to make for faster startup next time
      if(!read_only) ::memset(mmap_region + DRIVE_JOURNAL_START, 0, block_entries - (mmap_region + DRIVE_JOURNAL_START));

      // mark the disk as unmounted
      if(!read_only){
        assert(*reinterpret_cast<const uint32_t *>(mmap_region + 8) == 1);
        *reinterpret_cast<uint32_t *>(mmap_region + 8) = 0;
        assert(*reinterpret_cast<const uint32_t *>(mmap_region + 8) == 0);
      }

      // write out those last changes
      if((!read_only) && fdatasync(mmap_fd) && fdatasync(mmap_fd) && fdatasync(mmap_fd)) goto fatal;

      // check that we didn't corrupt the beginning of the drive
      assert(0xFFC0 == (((reinterpret_cast<uint64_t *>(block_entries))[0 * 2]) & 0xFFC0));
      assert(2 == ((((reinterpret_cast<uint64_t *>(block_entries))[0 * 2]) >> 16) & 0xFFFF));
      assert(0xFFC0 == (((reinterpret_cast<uint64_t *>(block_entries))[1 * 2]) & 0xFFC0));
      assert(0 == ((((reinterpret_cast<uint64_t *>(block_entries))[1 * 2]) >> 16) & 0xFFFF));

      if((mmap_region != MAP_FAILED) && (munmap(mmap_region, file_size)) && (munmap(mmap_region, file_size))) goto fatal;
    }

    ::munmap(block_indexes[0].data, block_indexes[0].data_length * Constants::NUM_SHARD_PARTITIONS);
    ::munmap(string_intern_lookup.next_available_counts, STRING_INTERN_COUNT_MEMORY_LENGTH);

    if((sync_fd != -1) && (close(sync_fd) != 0) && (close(sync_fd) != 0)) goto fatal;
    if((mmap_fd != -1) && (close(mmap_fd) != 0) && (close(mmap_fd) != 0)) goto fatal;
    sync_fd = -1;
    mmap_fd = -1;
    mmap_region = reinterpret_cast<char*>(MAP_FAILED);
    num_free_blocks = 0;
    next_block_to_use = 0;
    return;

  fatal:
    // TODO: possibly handle this better
    fatal_error("Fatal Error while unmounting a drive");
  }
}