server.cpp

#include <cstdlib>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <time.h>
#include <cstring>
#include <iostream>
#include "server.hpp"
#include "pool_allocator.hpp"
#include "message_dispatch.hpp"
#include "drive.hpp"
#include "logic_error.hpp"
#include "parser.hpp"

/*
PROTOCOL:
  INITIAL HANDSHAKE:
    Client:
      4 bytes: magic number (0x7EED9362)
      4 bytes: protocol version (1)
      56 bytes: reserved (0)

    Server:
      4 bytes: magic number (0x8F695F01)
      4 bytes: protocol version (1)
      56 bytes: reserved (0)

  REQUESTS:
    CLIENT -> SERVER WRITE TABLE REQUEST:
      1 byte: operation (0x1)
      3 bytes: request id
      10 bits: table index
      22 bits: table definition length
      n bytes: table definition

    CLIENT -> SERVER WRITE DATA REQUEST:
      1 byte: operation (0x2)
      3 bytes: request id
      10 bits: data type
      22 bits: data length
      n bytes: data

    CLIENT -> SERVER FENCE:
      1 byte: operation (special = 20x0)
      2 bytes: 0x0001 (begin transaction)
      1 byte: reserved (0)

    CLIENT -> SERVER FENCE:
      1 byte: operation (special = 20x0)
      2 bytes: 0x0002 (end transaction)
      1 byte: reserved (0)

    CLIENT -> SERVER READ ALL TABLES REQUEST:
      1 byte: operation (0x21)
      3 bytes: request id

    CLIENT -> SERVER SIMPLE QUERY:
      see simple_query.hpp

  
  RESPONSES:
    SERVER -> CLIENT WRITE ACKNOWLEDGEMENT OF SUCCESS
      1 byte: operation code (the initial code plus 0x80)
      3 bytes: request id

    SERVER -> CLIENT WRITE ACKNOWLEDGEMENT OF FAILURE
      1 byte: operation code (the initial code plus 0xC0)
      3 bytes: request id
      2 bytes: error code
      2 bytes: error message length
      n bytes: error message string

    SERVER -> CLIENT READ RESPONSE
      1 byte: operation code (the initial code plus 0xC0)
      3 bytes: request id
      4 bytes: response length
      n bytes: response data

    SERVER -> CLIENT READ RESPONSE FAILURE
      1 byte: operation code (the initial code plus 0xC0)
      3 bytes: request id
      2 bytes: error code
      2 bytes: error message length
      n bytes: error message string
*/
namespace AppenDB {

  // user_data setup 64 bits:
  // request_id = 24 bits
  // file descriptor = 16 bits
  // operation 8 bits
  // id 16 bits

    Server::Server(PortType p, uint32_t id, const char *ip_address) : _port(p), server_id(id) {
      // make a handshake for sending and recieving purposes
      uint32_t *hs = reinterpret_cast<uint32_t*>(handshake);
      uint32_t *hs_client = reinterpret_cast<uint32_t*>(handshake_client);
      hs[0] = MAGIC_NUMBER_SEND;
      hs[1] = 0x1;
      hs_client[0] = MAGIC_NUMBER;
      hs_client[1] = 0x1;
      for(int i = 2; i < 16; i++){
        hs[i] = 0;
        hs_client[i] = 0;
      }

      // setup the address and port to listen on
      sockaddr_in serv_addr;
      ::std::memset(&serv_addr, 0, sizeof(serv_addr));
      serv_addr.sin_family = AF_INET;
      serv_addr.sin_addr.s_addr = ip_address ? inet_addr(ip_address) : INADDR_ANY;
      serv_addr.sin_port = htons(_port);

      // create the server socket
      _server_socket_fd = ::socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);
      if(_server_socket_fd == -1){
        if((errno == EMFILE) || (errno == ENFILE) || (errno == ENOBUFS) || (errno == ENOMEM)){
          // TODO: bummer ... we are out of resources. This is likely very, very bad
        }
        fatal_error("Unable to create a server socket");
      }
      // set server socket to listen for new connections
      int yes = 1;
      if(-1 == ::setsockopt(_server_socket_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(int))) fatal_error("Unable to set server socket options");
      if(-1 == ::bind(_server_socket_fd, (struct sockaddr *) &serv_addr, sizeof(serv_addr))) fatal_error("Unable to bind server socket");
      if(-1 == ::listen(_server_socket_fd, MAX_CONNECTION_BACKLOG))  fatal_error("Unable to listen with the server socket");

      // setup the epoll file
      _epoll_fd = epoll_create1(0);
      if(-1 == _epoll_fd) fatal_error("Unable to setup epoll file descriptor");
      _epv.events = EPOLLIN | EPOLLPRI | EPOLLRDHUP | EPOLLHUP | EPOLLERR;
      _epv.data.ptr = 0;
      // add the server socket to the epoll list of sockets to poll
      int eprv = epoll_ctl(_epoll_fd, EPOLL_CTL_ADD, _server_socket_fd, &_epv);
      if(eprv == -1) fatal_error("Unable to attach the server socket to the epoll file descriptor");
    }

    Server::~Server() {
      // TODO: properly close socket and epoll when we stop server
      //if(::close(_server_socket_fd) && ::close(_server_socket_fd) && ::close(_server_socket_fd)) fatal_error("Error closing the server socket");
      //if(::close(_epoll_fd) && ::close(_epoll_fd) && ::close(_epoll_fd)) fatal_error("Error closing the epoll file descriptor");
      for(auto iter = _cd_by_fd.begin(); iter != _cd_by_fd.end(); ++iter){
        delete iter->second;
      }
    }

    void Server::operator()() {
      if(Dispatcher::running){
        ::std::cout << "Server running on port " << port() << ::std::endl;
        run_loop();
      }
      else {
        ::std::cout << "Server couldn't be started" << ::std::endl;
      }
    }

    void Server::create_server_connection(const char *ip_address, uint32_t port) {
      // setup the address and port to listen on
      sockaddr_in addr_info;
      ::std::memset(&addr_info, 0, sizeof(addr_info));
      addr_info.sin_family = AF_INET;
      addr_info.sin_addr.s_addr = inet_addr(ip_address);
      addr_info.sin_port = htons(port);

      // create the server socket
      int32_t fd = ::socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);
      if(fd < -1) {
        if((errno == EMFILE) || (errno == ENFILE) || (errno == ENOBUFS) || (errno == ENOMEM)){
          // TODO: bummer ... we are out of resources. This is likely very, very bad
        }
        fatal_error("Unable to create a client socket");
      }
      Parser *cd = create_parser(fd, true);
      cd->state = Parser::SERVER_HANDSHAKE;
      // TODO: consider deduplicating the above code and the code in the constructor with something like:
      //    uint32_t create_socket(in_addr_t ip_address, in_port_t, auto cb);
      //    then calling it with:
      // create_socket(inet_addr(ip_address), htons(port), [](uint32_t fd, sockaddr addr_info) {});
      if (::connect(fd, (struct sockaddr *) &addr_info, sizeof(addr_info)) < 0 && errno != EINPROGRESS) {
        // TODO: Probably not necessary for fatal error...
        fatal_error("Unable to connect to other server");
      }
    }

    void Server::check_timeout(){
      time_t now = time(0);
      for(auto iter = _cd_by_fd.begin(); iter != _cd_by_fd.end();){
        double duration = difftime(now, iter->second->timer);
        if(TIMEOUT_SECONDS < duration){
          Parser *p = iter->second;
          iter++;
          free_parser(p);
        }
        else{
          Parser *p = iter->second;
          char *data = "";
          iter++;
          if (p->send_buffer.offset_end != p->send_buffer.offset_start) {
            write_to_socket(data, 0, p->fd, p->id);
          }
        }
      }
    }

    void Server::accept_sockets(){
      for(;;){
        int new_socket_fd = ::accept4(_server_socket_fd, 0, 0, SOCK_NONBLOCK);
        if(-1 == new_socket_fd){
          switch(errno){
            // no more connections available to be established
            case EWOULDBLOCK: {
              return;
            } break;

            // logic errors
            case EOPNOTSUPP:
            case ENOTSOCK:
            case EINVAL:
            case EBADF:
            case EFAULT: {
              // these shouldn't ever happen
               fatal_error("Logic error while accepting connections");
            } break;

            // bad new connection, skip it
            case EPERM:
            case EPROTO:
            case ECONNABORTED:
            case EINTR:
            default: {
              continue;
            } break;

            // resource limits hit ... very bad
            case EMFILE:
            case ENFILE:
            case ENOBUFS:
            case ENOMEM: {
              // this is super bad ... there isn't a good way to handle this
              return;
            } break;
          }
        } else {
          create_parser(new_socket_fd, false);
        }
      }
    }

    // TODO: create structures that hold onto pending writes so that they can be completed even if more data is needed to be written later
    bool Server::write_to_socket(char *message, int length, int fd, uint16_t id){
      // TODO: obtain a lock on this Parser return false if lock not obtained
      Parser *cd = get_parser(fd, id);
      if (!cd) return true;
      cd->timer = time(0);
      if (cd->send_buffer.offset_start != cd->send_buffer.offset_end) {
        uint32_t send_queue_length = static_cast<uint32_t>(cd->send_buffer.offset_end - cd->send_buffer.offset_start);
        iovec iov[2];
        iov[0].iov_base = cd->send_buffer.buffer->buffer() + cd->send_buffer.offset_start;
        iov[0].iov_len = send_queue_length;
        iov[1].iov_base = message;
        iov[1].iov_len = length;
        ssize_t written = writev(fd, iov, 2);
        if (written < 0) {
          if (errno != EAGAIN) {
            free_parser(cd);
            return true;
          } else {
            written = 0;
          }
        }
        if (written < send_queue_length) {
          if (cd->send_buffer.offset_end + length < Constants::SEND_BUFFER_SIZE) {
            ::std::copy(message, message + length, cd->send_buffer.buffer->buffer() + cd->send_buffer.offset_end);
            cd->send_buffer.offset_start += written;
            cd->send_buffer.offset_end += length;
          } else {
            free_parser(cd);
          }
        } else if (written != (send_queue_length + length)) {
          cd->send_buffer.offset_start = 0;
          cd->send_buffer.offset_end = length - (written - send_queue_length);
          ::std::copy(message + (written - send_queue_length), message + length, cd->send_buffer.buffer->buffer());
        } else {
          cd->send_buffer.offset_start = 0;
          cd->send_buffer.offset_end = 0;
          cd->send_buffer.buffer->remove_ref();
          cd->send_buffer.has_ref = false;
        }
      } else {
        // MSG_NOSIGNAL prevents broken pipe from crashing program
        int written = write(fd, message, length);
        if (written < 0) {
          if (errno != EAGAIN) {
            free_parser(cd);
            return true;
          } else {
            written = 0;
          }
        }
        if (written != length) {
          PoolAllocator::RefCountedBuffer *new_buf = Dispatcher::sock_send_pool.allocate(server_id);
          uint32_t amount_left = length - written;
          if (new_buf && amount_left < Constants::SEND_BUFFER_SIZE) {
            new_buf->add_ref();
            cd->send_buffer.buffer = new_buf;
            cd->send_buffer.has_ref = true;
            cd->send_buffer.offset_start = 0;
            cd->send_buffer.offset_end = amount_left;
            ::std::copy(message + written, message + length, new_buf->buffer());
          } else {
            free_parser(cd);
          }
        }
      }
      return true;
    }

    Parser *Server::create_parser(uint32_t fd, bool use_out) {
      Parser *cd = new Parser(fd, (client_id << Constants::NUM_SOCKET_BITS) | server_id);
      client_id = static_cast<uint16_t>((client_id + 1) << Constants::NUM_SOCKET_BITS) >> Constants::NUM_SOCKET_BITS;
      if(client_id == 0) client_id++;
      _cd_by_fd[fd] = cd;
      epoll_event epv;
      epv.events = EPOLLIN | EPOLLPRI | EPOLLRDHUP | EPOLLHUP | EPOLLERR | (use_out ? EPOLLOUT : 0);
      epv.data.u64 = static_cast<uint64_t>(fd) + (static_cast<uint64_t>(cd->id) << 32);
      int eprv = epoll_ctl(_epoll_fd, EPOLL_CTL_ADD, fd, &epv);
      if(eprv == -1)  fatal_error("Unable to add a socket into the epoll list");
      return cd;
    }
    
    Parser *Server::get_parser(uint32_t fd, uint16_t id) {
      auto iter = _cd_by_fd.find(fd);
      if (iter == _cd_by_fd.end()) return nullptr;
      Parser *cd = iter->second;
      if(cd->id != id) return nullptr;
      return cd;
    }

    void Server::free_parser(Parser *p) {
      if (p->is_parsing) {
        p->state = ::AppenDB::Parser::CLEANUP;
      } else {
        _cd_by_fd.erase(p->fd);
        delete p;
      }
    }

/*
  TODO: lock ranges of bytes that are to be written out
    - once they are all written or the sockets are destroyed, we will know that we're able to stream
*/
    void Server::socket_read(Parser *cd){
      if (cd->receive_buffer.buffer == nullptr || Constants::MIN_READ_BUFFER_READ_SIZE > Constants::READ_BUFFER_SIZE - cd->receive_buffer.offset_end || reinterpret_cast<char *>(cd) != cd->receive_buffer.buffer->weak) {
        PoolAllocator::RefCountedBuffer *new_buffer = Dispatcher::sock_pool.allocate(server_id);
        if (new_buffer == nullptr) {
          const uint32_t cd_fd = cd->fd;
          const uint16_t cd_id = cd->id;
          process_append_events();
          cd = get_parser(cd_fd, cd_id);
          if (cd) socket_read(cd);
          return;
        }
        cd->set_buffer(new_buffer);
      }
      int amount = cd->read();
      // error with read
      if(amount < 0){
        switch(errno){
            case EWOULDBLOCK: {
              return;
            } break;

            // logic errors
            case EOPNOTSUPP:
            case EIO:
            case EISDIR:
            case ENOTSOCK:
            case EINVAL:
            case EBADF:
            case EFAULT: {
              // these shouldn't ever happen
              free_parser(cd);
              return;
            } break;

            // resource limits hit ... very bad
            case ENOBUFS:
            case ENOMEM:
            default:{
              // this is super bad ... there isn't a good way to handle this
               fatal_error("Out of memory or buffers");
              return;
            } break;
        }
      } else if (amount == 0) { // socket closed if fd is eof (due to socket being closed)
        free_parser(cd);
        return;
      } else {
        if (cd->state == ::AppenDB::Parser::PARSABLE) {
          // parse data and delete parser if necessary
          if (!cd->parse(this)) free_parser(cd);
        }
        else if (cd->state == ::AppenDB::Parser::HANDSHAKE && amount == 64 && memcmp(handshake_client, cd->receive_buffer.buffer->buffer() + cd->receive_buffer.offset_start, 64) == 0) {
          cd->state = ::AppenDB::Parser::PARSABLE;
          cd->receive_buffer.offset_start = 0;
          cd->receive_buffer.offset_end = 0;
          cd->receive_buffer.buffer->remove_ref();
          cd->receive_buffer.has_ref = false;
          write_to_socket(handshake, 64, cd->fd, cd->id);
        } else if (cd->state == ::AppenDB::Parser::SERVER_HANDSHAKE && amount == 64 && memcmp(handshake, cd->receive_buffer.buffer->buffer() + cd->receive_buffer.offset_start, 64) == 0) {
          cd->state = ::AppenDB::Parser::PARSABLE;
          cd->receive_buffer.offset_start = 0;
          cd->receive_buffer.offset_end = 0;
          cd->receive_buffer.buffer->remove_ref();
          cd->receive_buffer.has_ref = false;
        } else { // handshake was not fully recieved or is incorrect
          free_parser(cd);
        }
      }
    }

    void Server::allocate_op_memory(size_t num_spots, auto return_buffer) {
      if (!op_buffer.buffer || (num_spots + op_buffer.amount_used) > Constants::OP_BUFFER_MAX_OPS) {
        PoolAllocator::RefCountedBuffer *new_buf = nullptr;
        while (!(new_buf = Dispatcher::op_pool.allocate(server_id))) {
          this->process_append_events();
        }
        op_buffer.set_buffer(new_buf);
      }
      const uint8_t size_of_write_operation = 32;
      char *memory = op_buffer.buffer->buffer() + (op_buffer.amount_used * size_of_write_operation);
      op_buffer.amount_used += num_spots;
      return_buffer(memory, this->op_buffer.buffer);
    }

    void Server::allocate_row_memory(size_t size_needed, auto return_buffer) {
      if (!row_buffer.buffer || (size_needed + row_buffer.amount_used) > Constants::READ_BUFFER_SIZE) {
        PoolAllocator::RefCountedBuffer *new_buf = nullptr;
        while (!(new_buf = Dispatcher::sock_pool.allocate(server_id))) {
          this->process_append_events();
        }
        row_buffer.set_buffer(new_buf);
      }
      char *memory = row_buffer.buffer->buffer() + row_buffer.amount_used;
      row_buffer.amount_used += size_needed;
      return_buffer(memory, this->row_buffer.buffer);
    }

    void Server::socket_ready(Parser *cd, uint32_t events){
      if(events & (EPOLLRDHUP | EPOLLHUP | EPOLLERR)){
        free_parser(cd);
        return;
      }
      cd->timer = time(0);
      // read from socket into global buffer for parsing
      socket_read(cd);
    }

    void Server::process_epoll_events(int nresults){
      for(int i = 0; i < nresults; i++){
        if(!Dispatcher::running) return;
        epoll_event &e = _epoll_events[i];

        // This is server socket that needs a handshake
        if (e.events & EPOLLOUT) {
          Parser *cd = get_parser(e.data.u64 & 0xFFFFFFFF, (e.data.u64 >> 32) & 0xFFFF);
          if (!cd) continue;
          if (cd->state != Parser::SERVER_HANDSHAKE) emit_logic_error("Bad handshake");
          e.events = EPOLLIN | EPOLLPRI | EPOLLRDHUP | EPOLLHUP | EPOLLERR;
          epoll_ctl(_epoll_fd, EPOLL_CTL_MOD, cd->fd, &e);
          write_to_socket(handshake_client, 64, cd->fd, cd->id);
        }
        // this is a client socket
        else if(e.data.u64) {
          Parser *cd = get_parser(e.data.u64 & 0xFFFFFFFF, (e.data.u64 >> 32) & 0xFFFF);
          if (!cd) continue;
          // TODO: Need to obtain a lock on this Parser before calling socket_ready
          socket_ready(cd, e.events);
        }

        // this is the server socket
        else {
          accept_sockets();
        }
      }
      // TODO: this can cause issues if items are deleted.
      check_timeout();
    }

    void Server::process_append_events() {
      uint64_t sock_op;
      while ((sock_op = Dispatcher::socket_queue->pop())) {
        uint8_t op = (sock_op >> 56) & 0xFF;
        if (op == Parser::READ_COUNT_REQUEST || op == Parser::READ_SHARD_REQUEST) {
          ReadOperationBase &read_op = *reinterpret_cast<ReadOperationBase*>(sock_op & ((1ULL << 48) - 1));
          read_op.finalize();
        } else {
          uint32_t offset = sock_op & 0xFFFFFFFF;
          uint16_t response_code = (sock_op >> 36) & 0xFFFF;
          WriteOperation &write_op = *reinterpret_cast<WriteOperation *>(Dispatcher::op_pool.buffer() + (sizeof(WriteOperation) * offset));
          if (response_code == 1) {
            if (write_op.id() == 0) {
              uint32_t shard = write_op.fd();
              Drive::get_any_drive().string_intern_lookup().finish_operation(shard);
              assert(shard < Constants::DEFAULT_NUM_SHARDS);
              write_op.finalize();
            } else {
              uint32_t message = (write_op.rid() << 8) + write_op.op() + 0x80;
              char *output = reinterpret_cast<char*>(&message);
              if (write_to_socket(output, 4, write_op.fd(), write_op.id())) {
                write_op.finalize();
              } else {
                // Couldn't obtain the socket try again in a bit...
                // TODO: handle potential to drop op if there is now no more room on the queue...
                bool r = Dispatcher::socket_queue->push(sock_op);
                assert(r);
              }
            }
          } else {
            // TODO: handle rollback messages
            write_op.finalize();
          }
        }
      }
    }

    void Server::run_loop() {
      while (Dispatcher::running) {
        int nresults = epoll_wait(_epoll_fd, _epoll_events, MAX_EPOLL_EVENTS, 2);
        process_epoll_events(nresults);
        process_append_events();
        if (nresults == 0) {
          for(auto iter = _cd_by_fd.begin(); iter != _cd_by_fd.end();){
            Parser *cd = iter->second;
            ++iter;
            if (cd->receive_buffer.has_ref && cd->receive_buffer.offset_start != cd->receive_buffer.offset_end) {
              if (!cd->parse(this)) free_parser(cd);
            }
          }
        }
      }
    }

    void Server::process_outstanding_block_forever(){
      int nresults = epoll_wait(_epoll_fd, _epoll_events, MAX_EPOLL_EVENTS, -1);
      process_epoll_events(nresults);
    }

    void Server::process_outstanding_block_with_timeout(int ms_timeout){
      int nresults = epoll_wait(_epoll_fd, _epoll_events, MAX_EPOLL_EVENTS, ms_timeout);
      process_epoll_events(nresults);
    }

    void Server::process_outstanding_non_block(){
      int nresults = epoll_wait(_epoll_fd, _epoll_events, MAX_EPOLL_EVENTS, 0);
      process_epoll_events(nresults);
    }

}