main.cpp

#ifndef EXTERNAL_TOOL
#include "drive.hpp"
#include "server.hpp"
#include "shard_worker.hpp"
#include "test/generate_data.hpp"
#include "constants.hpp"
#include <iostream>
#include <cstring>
#include <signal.h>
#include <cstdlib>
#include <thread>
#include <vector>

::AppenDB::Drive drive;
::std::vector<::std::thread> threads;

void format(const char *name){
  bool is_formatted = ::AppenDB::Drive::format(name);
  ::std::cout << "is_formatted: " << is_formatted << ::std::endl;
}

void attach(const char *name, bool force = false){
  drive.attach(name, force);
  if(drive){
    ::std::cout << name << " mounted" << ::std::endl;
  }
  else {
    ::std::cout << name << " not mounted" << ::std::endl;
  }
}

void detach(){
  drive.detach();
}

void write_read_test(){
  ::AppenDB::Test::test_key_value_data_box(drive, 1, 900, 300, 303);
  ::AppenDB::Test::test_key_value_data_box(drive, 904, 1025, 300, 23687);
  ::AppenDB::Test::test_key_value_data_box(drive, 1026, 2000, 300, 173303);
  ::AppenDB::Test::test_key_value_data_box(drive, 2001, 2003, 300, 23687);
  ::AppenDB::Dispatcher::running = false;
}

int print_help(const char* exe_name){
  ::std::cerr << "Expected to be called: \n$ " << exe_name << " <block device> <command>\n\nCommands accepted:\n  format: format the drive (all data lost)"
   << "\n  check: make sure drive is valid\n  cleanup: If another process died that was attached to the drive, this can be run to force the mount flag to be cleared"
   << "\n  test: run tests on the drive\n  server: start a server\n\nExample:\n$ " << exe_name << " /dev/loop4 format" << ::std::endl;
  return 1;
}

static void on_sigterm(int sig, siginfo_t *si, void *unused) {
  ::std::cout << "Termination signal received. Stopping server" << ::std::endl;
  ::AppenDB::Dispatcher::running = false;
}

void setup_signal_handler(){
  struct sigaction sa;
  sa.sa_flags = SA_SIGINFO;
  sigemptyset(&sa.sa_mask);
  sa.sa_sigaction = on_sigterm;
  signal(SIGPIPE, SIG_IGN);
  assert(0 == ::sigaction(SIGINT, &sa, 0));
  assert(0 == ::sigaction(SIGTERM, &sa, 0));
  assert(0 == ::sigaction(SIGQUIT, &sa, 0));
  //assert(0 == ::sigaction(SIGPIPE, &sa, 0));
  assert(0 == ::sigaction(SIGALRM, &sa, 0));
  assert(0 == ::sigaction(SIGUSR1, &sa, 0));
  assert(0 == ::sigaction(SIGUSR2, &sa, 0));
}

void spawn_shard_threads() {
  const size_t NUM_SHARD_THREADS = ::AppenDB::Constants::NUM_SHARD_PARTITIONS;
  for (size_t i = 0; i < NUM_SHARD_THREADS; ++i) {
    threads.push_back(::std::thread(::AppenDB::ShardWorker(), i));
  }
}

void join_threads() {
  for (auto &t : threads) t.join();
}

void spawn_tests(const char* drive_to_use) {
  attach(drive_to_use);
  ::AppenDB::Dispatcher::running = true;
  setup_signal_handler();
  spawn_shard_threads();
  threads.push_back(::std::thread(write_read_test));
  join_threads();
  detach();
}

void spawn_server(const char* drive_to_use, const char* ip_address) {
  attach(drive_to_use);
  ::AppenDB::Dispatcher::running = true;
  setup_signal_handler();
  spawn_shard_threads();
  // TODO: Don't use functors for these since it will end up creating more than one and can cause problems
  //  if there are any side effects with file descriptors
  threads.push_back(::std::thread(::AppenDB::Server(64776, 0, ip_address)));
  join_threads();
  detach();
}

int main(int argc, const char *argv[]){
  if(argc < 3) return print_help(argv[0]);

  const char* drive_to_use = argv[1];
  const char* command = argv[2];

  if(!::std::strcmp(command, "truncate_block")){
    if(argc != 5) return print_help(argv[0]);
    uint32_t const block_id = ::std::atoi(argv[3]);
    uint32_t const new_block_size = ::std::atoi(argv[4]);
    attach(drive_to_use);
    drive.repair_truncate_block(block_id, new_block_size);
    ::std::cout << "Block #" << block_id << " truncated to " << new_block_size << " bytes" << ::std::endl;
    detach();
  }

  if(argc < 3) return print_help(argv[0]);

  if(!::std::strcmp(command, "format")) format(drive_to_use);
  else if(!::std::strcmp(command, "check")){
    attach(drive_to_use);
    detach();
  }
  else if(!::std::strcmp(command, "cleanup")){
    attach(drive_to_use, true);
    detach();
  }
  else if(!::std::strcmp(command, "test")){
    spawn_tests(drive_to_use);
  }
  else if(!::std::strcmp(command, "server")){
    spawn_server(drive_to_use, argc > 3 ? argv[3] : nullptr);
  }
  else return print_help(argv[0]);

  return 0;
}
#endif