server.cpp

#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <poll.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

#include <map>
#include <string>
#include <vector>

#include "hashtable.h"

#define container_of(ptr, T, member)                   \
    ({                                                 \
        const typeof(((T*)0)->member)* __mptr = (ptr); \
        (T*)((char*)__mptr - offsetof(T, member));     \
    })

static void msg(const char* msg) { fprintf(stderr, "%s\n", msg); }

static void msg_errno(const char* msg) {
    fprintf(stderr, "[errno:%d] %s\n", errno, msg);
}

static void die(const char* msg) {
    int err = errno;
    fprintf(stderr, "[%d] %s", err, msg);
    abort();
}

static void fd_set_nb(int fd) {
    errno = 0;
    int flags = fcntl(fd, F_GETFL, 0);
    if (errno) {
        die("fcntl error");
        return;
    }

    flags |= O_NONBLOCK;

    errno = 0;
    (void)fcntl(fd, F_SETFL, flags);
    if (errno) {
        die("fcntl error");
    }
}

const size_t k_max_msg = 32 << 20;  // larger than kernel buffer

struct Conn {
    int fd = -1;
    bool want_read = false;
    bool want_write = false;
    bool want_close = false;
    // data collected from non blocking read
    // i.e., data to be parsed by the application
    std::vector<uint8_t> incoming;
    // data that will be used for non blocking write
    // i.e. responses generated by application
    std::vector<uint8_t> outgoing;
};

// append to back
static void buf_append(std::vector<uint8_t>& buf, const uint8_t* data,
                       size_t n) {
    buf.insert(buf.end(), data, data + n);
}

// remove from front
static void buf_consume(std::vector<uint8_t>& buf, size_t n) {
    buf.erase(buf.begin(), buf.begin() + n);
}

static Conn* handle_accept(int fd) {
    struct sockaddr_in client_addr = {};
    socklen_t addrlen = sizeof(client_addr);
    int connfd = accept(fd, (struct sockaddr*)&client_addr, &addrlen);
    if (connfd < 0) {
        msg_errno("accept() error");
        return NULL;
    }

    uint32_t ip = client_addr.sin_addr.s_addr;
    fprintf(stderr, "new client from %u.%u.%u.%u:%u\n", ip & 255,
            (ip >> 8) & 255, (ip >> 16) & 255, ip >> 24,
            ntohs(client_addr.sin_port));

    // set connfd as non blocking
    fd_set_nb(connfd);

    // create new `struct Conn`
    Conn* conn = new Conn();
    conn->fd = connfd;
    conn->want_read = true;
    return conn;
}

const size_t k_max_args = 200 * 1000;

static bool read_u32(const uint8_t*& cur, const uint8_t* end, uint32_t& out) {
    // cur is reference to pointer, updating this pointer changes the original
    // pointer as well
    if (cur + 4 > end) {
        return false;
    }
    memcpy(&out, cur, 4);
    cur += 4;
    return true;
}

static bool read_str(const uint8_t*& cur, const uint8_t* end, size_t n,
                     std::string& out) {
    if (cur + n > end) {
        return false;
    }
    out.assign(cur, cur + n);
    cur += n;
    return true;
}

// +------+-----+------+-----+------+-----+-----+------+
// | nstr | len | str1 | len | str2 | ... | len | strn |
// +------+-----+------+-----+------+-----+-----+------+
static int32_t parse_req(const uint8_t* data, size_t size,
                         std::vector<std::string>& out) {
    const uint8_t* end = data + size;
    uint32_t n_str = 0;
    if (!read_u32(data, end, n_str)) {
        return -1;
    }
    if (n_str > k_max_args) {
        return -1;
    }

    while (out.size() < n_str) {
        uint32_t len = 0;
        if (!read_u32(data, end, len)) {
            return -1;
        }
        out.push_back(std::string());
        if (!read_str(data, end, len, out.back())) {
            return -1;
        }
    }

    if (data != end) {
        return -1;  // trailing garbage
    }
    return 0;
}

// Response::status
enum {
    RES_OK = 0,
    RES_ERR = 1,  // error
    RES_NX = 2,   // key not found
};

struct Response {
    uint32_t status = 0;
    std::vector<uint8_t> data;
};

// use custom hashmap instead of STL for scalability, low latency
static struct {
    HMap db;
} g_data;

struct Entry {
    HNode node;  // hashtable node
    std::string key;
    std::string val;
};

static bool entry_eq(HNode* lhs, HNode* rhs) {
    struct Entry* le = container_of(lhs, struct Entry, node);
    struct Entry* re = container_of(rhs, struct Entry, node);
    return le->key == re->key;
}

// FNV hash
static uint64_t str_hash(const uint8_t* data, size_t len) {
    uint64_t h = 0xcbf29ce484222325;  // FNV_offset_basis
    for (size_t i = 0; i < len; i++) {
        h = h * 0x00000100000001b3;  // FNV_prime
        h ^= data[i];
    }
    return h;
}

static void do_get(std::vector<std::string>& cmd, Response& out) {
    // Entry is dummy container used just for lookup
    Entry key;
    key.key.swap(cmd[1]);
    key.node.hcode = str_hash((uint8_t*)key.key.data(), key.key.size());
    // hashtable lookup
    HNode* node = hm_lookup(&g_data.db, &key.node, &entry_eq);
    if (!node) {
        out.status = RES_NX;
        return;
    }
    // get the value from node
    const std::string& val = container_of(node, Entry, node)->val;
    assert(val.size() <= k_max_msg);
    out.data.assign(val.begin(), val.end());
}

static void do_set(std::vector<std::string>& cmd, Response&) {
    Entry key;
    key.key.swap(cmd[1]);
    key.node.hcode = str_hash((uint8_t*)key.key.data(), key.key.size());
    // hashtable lookup
    HNode* node = hm_lookup(&g_data.db, &key.node, &entry_eq);
    if (node) {
        container_of(node, Entry, node)->val.swap(cmd[2]);
    } else {
        // key not found then create new entry
        Entry* ent = new Entry();
        ent->key.swap(key.key);
        ent->val.swap(cmd[2]);
        ent->node.hcode = key.node.hcode;
        hm_insert(&g_data.db, &ent->node);
    }
}

static void do_delete(std::vector<std::string>& cmd, Response&) {
    Entry key;
    key.key.swap(cmd[1]);
    key.node.hcode = str_hash((uint8_t*)key.key.data(), key.key.size());
    HNode* node = hm_delete(&g_data.db, &key.node, &entry_eq);
    if (node) {
        // deallocate the KV pair container.
        delete container_of(node, Entry, node);
    }
}

static void do_request(std::vector<std::string>& cmd, Response& out) {
    if (cmd.size() == 2 && cmd[0] == "get") {
        return do_get(cmd, out);
    } else if (cmd.size() == 3 && cmd[0] == "set") {
        return do_set(cmd, out);
    } else if (cmd.size() == 2 && cmd[0] == "del") {
        return do_delete(cmd, out);
    } else {
        out.status = RES_ERR;
    }
}

static void make_response(const Response& resp, std::vector<uint8_t>& out) {
    // 4 bytes for the status code
    uint32_t resp_len = 4 + (uint32_t)resp.data.size();
    buf_append(out, (const uint8_t*)&resp_len, 4);
    buf_append(out, (const uint8_t*)&resp.status, 4);
    buf_append(out, resp.data.data(), resp.data.size());
}

// process 1 request if there is enough data
static bool try_one_request(Conn* conn) {
    // try to parse the messsage header(msg length)
    if (conn->incoming.size() < 4) {
        return false;  // needs more read
    }

    uint32_t len = 0;
    memcpy(&len, conn->incoming.data(), 4);
    if (len > k_max_msg) {
        msg("too long");
        conn->want_close = true;
        return false;
    }

    // msg body
    if (4 + len > conn->incoming.size()) {
        return false;  // needs more read
    }

    const uint8_t* request = &conn->incoming[4];

    // do some application logic for the request
    std::vector<std::string> cmd;
    if (parse_req(request, len, cmd) < 0) {
        msg("bad request");
        conn->want_close = true;
        return false;
    }

    Response resp;
    do_request(cmd, resp);
    make_response(resp, conn->outgoing);

    // remove the request message once request is processed
    buf_consume(conn->incoming, 4 + len);
    return true;
}

// application callback when the socket is writable, i.e. non blocking write
static void handle_write(Conn* conn) {
    assert(conn->outgoing.size() > 0);
    ssize_t rv = write(conn->fd, &conn->outgoing[0], conn->outgoing.size());
    if (rv < 0 && errno == EAGAIN) {
        return;  // not ready
    }

    if (rv < 0) {
        msg_errno("write() error");
        conn->want_close = true;
        return;
    }

    // remove written data from outgoing
    buf_consume(conn->outgoing, (size_t)rv);

    // if no other data to be written, read next request
    if (conn->outgoing.size() == 0) {
        conn->want_read = true;
        conn->want_write = false;
    }
    // else want_write=true, response for current request not fully written yet.
}

// application callback when the socket is readable
static void handle_read(Conn* conn) {
    // read some data
    uint8_t buf[64 * 1024];
    ssize_t rv = read(conn->fd, buf, sizeof(buf));
    if (rv < 0 && errno == EAGAIN) {
        return;  // not ready for read yet
    }
    if (rv < 0) {
        msg_errno("read() error");
        conn->want_close = true;
        return;
    }

    // handle EOF
    if (rv == 0) {
        if (conn->incoming.size() == 0) {
            msg("client closed");
        } else {
            msg("unexpected EOF");
        }
        conn->want_close = true;
        return;
    }

    // write the data read to incoming
    buf_append(conn->incoming, buf, (size_t)rv);

    // parse requests and generate responses
    while (try_one_request(conn)) {
    }  // support piplelining

    // if outgoing has response, change readiness to write
    if (conn->outgoing.size() > 0) {
        conn->want_read = false;
        conn->want_write = true;
        // socket is likely ready for write
        // try to write without waiting for next iteration
        return handle_write(conn);
    }
    // else need read
}

int main() {
    int fd = socket(AF_INET, SOCK_STREAM, 0);  // IPV4, TCP
    if (fd == -1) {
        die("socket()");
    }
    int val = 1;
    // for reusing same ip:port after restart
    setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &val, sizeof(val));

    // bind socket to ip:port = 0.0.0.0:1234
    struct sockaddr_in addr = {};
    addr.sin_family = AF_INET;
    addr.sin_port = htons(1234);      // unsigned host to network short
    addr.sin_addr.s_addr = htonl(0);  // unsigned host to network long

    // address size is needed because addr size is different for ipv6 and ipv4
    int rv = bind(fd, (const sockaddr*)&addr, sizeof(addr));
    if (rv) {
        die("bind()");
    }

    // set the listen fd to non-blocking mode
    fd_set_nb(fd);

    // listen
    rv = listen(fd, SOMAXCONN);
    if (rv) {
        die("listen()");
    }

    // map of all client connections keyed by fd
    std::vector<Conn*> fd2conn;
    // event loop
    std::vector<struct pollfd> poll_args;
    while (true) {
        /* In the event loop, everytime we clear the poll_args
           and recreate the poll_args for the previous connections using fd2conn
        */
        // prepare arguments of the poll()
        poll_args.clear();
        // put listening socket in the first position
        struct pollfd pfd = {fd, POLLIN, 0};
        poll_args.push_back(pfd);

        // rest are connection sockets
        for (Conn* conn : fd2conn) {
            if (!conn) {
                continue;
            }
            // poll for error
            struct pollfd pfd = {conn->fd, POLLERR, 0};

            if (conn->want_read) {
                pfd.events |= POLLIN;
            }
            if (conn->want_write) {
                pfd.events |= POLLOUT;
            }
            poll_args.push_back(pfd);
        }

        int rv = poll(poll_args.data(), (nfds_t)poll_args.size(), -1);
        if (rv < 0 && errno == EINTR) {
            continue;  // interrupted by a signal
        }

        if (rv < 0) {
            die("poll error");
        }

        // handle listening socket
        if (poll_args[0].revents) {
            if (Conn* conn = handle_accept(fd)) {
                // resize if vector is smaller
                if (fd2conn.size() <= (size_t)conn->fd) {
                    fd2conn.resize(conn->fd + 1);
                }
                // same fd doesn't exist for different connection
                assert(!fd2conn[conn->fd]);
                // put it into the map
                fd2conn[conn->fd] = conn;
            }
        }

        // handle connection sockets
        for (size_t i = 1; i < poll_args.size(); i++) {
            uint32_t ready = poll_args[i].revents;
            if (ready == 0) {
                continue;
            }

            Conn* conn = fd2conn[poll_args[i].fd];
            if (ready & POLLIN) {
                assert(conn->want_read);
                handle_read(conn);  // application logic
            }
            if (ready & POLLOUT) {
                assert(conn->want_write);
                handle_write(conn);  // application logic
            }

            // close the socket if error or want_close
            if ((ready & POLLERR) || conn->want_close) {
                close(conn->fd);
                fd2conn[conn->fd] = NULL;
                delete conn;
            }
        }
    }
    return 0;
}