yrmcds のバイナリプロトコルの自動テストツール（書きかけ）

protocol_binary.cpp

#include "../src/memcache.hpp"

#include <cybozu/dynbuf.hpp>
#include <cybozu/tcp.hpp>
#define TEST_DISABLE_AUTO_RUN
#include <cybozu/test.hpp>
#include <cybozu/util.hpp>

#include <cstdint>
#include <cstring>
#include <fcntl.h>
#include <string>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#include <vector>

using yrmcds::memcache::binary_command;
using yrmcds::memcache::binary_status;
using yrmcds::memcache::item;

const char* g_server = nullptr;
std::uint16_t g_port = 11211;

typedef char opaque_t[4];
const std::size_t BINARY_HEADER_SIZE = 24;

int connect_server() {
    int s = cybozu::tcp_connect(g_server, g_port);
    if( s == -1 ) return -1;
    ::fcntl(s, F_SETFL, ::fcntl(s, F_GETFL, 0) & ~O_NONBLOCK);
    return s;
}


// compose binary request
class request {
public:
    request(binary_command cmd,
            std::uint16_t key_len, const char* key,
            char extras_len, const char* extras,
            std::uint32_t data_len, const char* data,
            opaque_t* opaque = nullptr, std::uint64_t cas = 0):
        m_data(BINARY_HEADER_SIZE + key_len + extras_len + data_len),
        m_p(&m_data[0])
    {
        m_p[0] = '\x80';
        m_p[1] = (char)cmd;
        cybozu::hton(key_len, m_p+2);
        m_p[4] = extras_len;
        m_p[5] = 0;
        m_p[6] = 0;
        m_p[7] = 0;
        std::uint32_t total_len = key_len + extras_len + data_len;
        cybozu::hton(total_len, m_p+8);
        if( opaque != nullptr )
            std::memcpy(m_p+12, opaque, sizeof(opaque_t));
        cybozu::hton(cas, m_p+16);
        std::memcpy(m_p, extras, extras_len);
        std::memcpy(m_p+BINARY_HEADER_SIZE+extras_len, key, key_len);
        std::memcpy(m_p+BINARY_HEADER_SIZE+extras_len+key_len, data, data_len);
    }

    const char* data() {
        return m_p;
    }

    std::size_t length() const noexcept {
        return m_data.size();
    }

private:
    std::vector<char> m_data;
    char* const m_p;
};


// parse binary response
class response {
public:
    response() {}
    response(response&&) noexcept = default;

    // Return length of the response.
    //
    // Return length of the response.
    // If the response is incomplete, zero is returned.
    std::size_t length() const noexcept { return m_response_len; }

    // Response status, if determined by the request.
    binary_status status() const noexcept { return m_status; }

    // Return `true` if the response status was not OK.
    bool error() {
        return m_status != binary_status::OK;
    }

    // Return the command type.
    binary_command command() const noexcept { return m_command; }

    // Return `key`.
    item key() const noexcept { return m_key; }

    // Return `opaque` sent with the request.
    const char* opaque() const noexcept { return m_p + 12; }

    // Return `cas unique` sent with CAS command.
    std::uint64_t cas_unique() const noexcept { return m_cas_unique; }

    // Return `flags` sent with storage commands.
    std::uint32_t flags() const noexcept { return m_flags; }

    // Return data block sent with storage commands.
    item data() const noexcept { return m_data; }

    // Return an unsigned 64bit integer value for increment or decrement.
    std::uint64_t value() const noexcept { return m_value; }

    bool parse(const char* p, std::size_t len);

private:
    const char* m_p = nullptr;
    std::size_t m_len = 0;
    std::size_t m_response_len = 0;
    binary_status m_status;
    binary_command m_command;
    item m_key;
    std::uint64_t m_cas_unique;
    std::uint32_t m_flags = 0;
    item m_data;
    std::uint64_t m_value = 0;
};

bool response::parse(const char* p, std::size_t len) {
    m_p = p;
    m_len = len;
    if( m_len < BINARY_HEADER_SIZE )
        return false;
    cybozu_assert( *m_p == '\x81' );
    std::uint32_t total_len;
    cybozu::ntoh(m_p + 8, total_len);
    if( m_len < (BINARY_HEADER_SIZE + total_len) )
        return false;
    m_response_len = BINARY_HEADER_SIZE + total_len;

    m_command = (binary_command)*(unsigned char*)(m_p + 1);
    std::uint16_t key_len;
    cybozu::ntoh(m_p + 2, key_len);
    std::uint8_t extras_len = *(unsigned char*)(m_p + 4);
    cybozu_assert( total_len >= (key_len + extras_len) );
    if( key_len > 0 )
        m_key = item(m_p + (BINARY_HEADER_SIZE + extras_len), key_len);

    std::uint16_t i_status;
    cybozu::ntoh(m_p + 6, i_status);
    m_status = (binary_status)i_status;
    cybozu::ntoh(m_p + 16, m_cas_unique);

    std::size_t data_len = total_len - key_len - extras_len;
    if( data_len > 0 )
        m_data = item(m_p + (BINARY_HEADER_SIZE + extras_len + key_len),
                      data_len);

    if( extras_len > 0 ) {
        cybozu_assert( extras_len == 4 );
        cybozu::ntoh(m_p + BINARY_HEADER_SIZE, m_flags);
    }
    return true;
}


class client {
public:
    client(): m_socket(connect_server()), m_buffer(1 << 20) {
        cybozu_assert( m_socket != -1 );
    }
    ~client() { ::close(m_socket); }

    bool get_response(response& resp) {
        m_buffer.erase(m_last_response_size);
        if( m_buffer.empty() ) {
            if( ! recv() ) {
                cybozu_assert( m_buffer.empty() );
                return false;
            }
        }
        while( true ) {
            if( resp.parse(m_buffer.data(), m_buffer.size()) ) {
                m_last_response_size = resp.length();
                return true;
            }
            cybozu_assert( recv() );
        }
    }

    void noop(opaque_t* opaque) {
        request req(binary_command::Noop,
                    0, nullptr, 0, nullptr, 0, nullptr, opaque);
        send(req.data(), req.length());
    }

    void get(const std::string& key, bool q) {
        request req(q ? binary_command::GetQ : binary_command::Get,
                    (std::uint16_t)key.size(), key.data(),
                    0, nullptr, 0, nullptr);
        send(req.data(), req.length());
    }

    void getk(const std::string& key, bool q) {
        request req(q ? binary_command::GetKQ : binary_command::GetK,
                    (std::uint16_t)key.size(), key.data(),
                    0, nullptr, 0, nullptr);
        send(req.data(), req.length());
    }

    void get_and_touch(const std::string& key, std::uint32_t expire, bool q) {
        char extra[4];
        cybozu::hton(expire, extra);
        request req(q ? binary_command::GaTQ : binary_command::GaT,
                    (std::uint16_t)key.size(), key.data(),
                    sizeof(extra), extra, 0, nullptr);
        send(req.data(), req.length());
    }

    void lock_and_get(const std::string& key, bool q) {
        request req(q ? binary_command::GaTQ : binary_command::GaT,
                    (std::uint16_t)key.size(), key.data(),
                    0, nullptr, 0, nullptr);
        send(req.data(), req.length());
    }

    void set(const std::string& key, const std::string& data, bool q,
             std::uint32_t flags, std::uint32_t expire, std::uint64_t cas = 0) {
        char extra[8];
        cybozu::hton(flags, extra);
        cybozu::hton(expire, &extra[4]);
        request req(q ? binary_command::SetQ : binary_command::Set,
                    (std::uint16_t)key.size(), key.data(),
                    sizeof(extra), extra,
                    (std::uint32_t)data.length(), data.data(),
                    nullptr, cas);
        send(req.data(), req.length());
    }

    void replace(const std::string& key, const std::string& data, bool q,
                 std::uint32_t flags, std::uint32_t expire, std::uint64_t cas = 0) {
        char extra[8];
        cybozu::hton(flags, extra);
        cybozu::hton(expire, &extra[4]);
        request req(q ? binary_command::ReplaceQ : binary_command::Replace,
                    (std::uint16_t)key.size(), key.data(),
                    sizeof(extra), extra,
                    (std::uint32_t)data.length(), data.data(),
                    nullptr, cas);
        send(req.data(), req.length());
    }

    void add(const std::string& key, const std::string& data, bool q,
             std::uint32_t flags, std::uint32_t expire, std::uint64_t cas = 0) {
        char extra[8];
        cybozu::hton(flags, extra);
        cybozu::hton(expire, &extra[4]);
        request req(q ? binary_command::AddQ : binary_command::Add,
                    (std::uint16_t)key.size(), key.data(),
                    sizeof(extra), extra,
                    (std::uint32_t)data.length(), data.data(),
                    nullptr, cas);
        send(req.data(), req.length());
    }

    void replace_and_unlock(const std::string& key, const std::string& data,
                            bool q, std::uint32_t flags, std::uint32_t expire,
                            std::uint64_t cas = 0) {
        char extra[8];
        cybozu::hton(flags, extra);
        cybozu::hton(expire, &extra[4]);
        request req(q ? binary_command::RaUQ : binary_command::RaU,
                    (std::uint16_t)key.size(), key.data(),
                    sizeof(extra), extra,
                    (std::uint32_t)data.length(), data.data(),
                    nullptr, cas);
        send(req.data(), req.length());
    }

    void incr(const std::string& key, std::uint64_t value, bool q,
              std::uint32_t expire = ~(std::uint32_t)0,
              std::uint64_t initial = 0) {
        char extra[20];
        cybozu::hton(value, extra);
        cybozu::hton(initial, &extra[8]);
        cybozu::hton(expire, &extra[16]);
        request req(q ? binary_command::IncrementQ : binary_command::Increment,
                    (std::uint16_t)key.size(), key.data(),
                    sizeof(extra), extra, 0, nullptr);
        send(req.data(), req.length());
    }

    void decr(const std::string& key, std::uint64_t value, bool q,
              std::uint32_t expire = ~(std::uint32_t)0,
              std::uint64_t initial = 0) {
        char extra[20];
        cybozu::hton(value, extra);
        cybozu::hton(initial, &extra[8]);
        cybozu::hton(expire, &extra[16]);
        request req(q ? binary_command::DecrementQ : binary_command::Decrement,
                    (std::uint16_t)key.size(), key.data(),
                    sizeof(extra), extra, 0, nullptr);
        send(req.data(), req.length());
    }

    void remove(const std::string& key, bool q) {
        request req(q ? binary_command::DeleteQ : binary_command::Delete,
                    (std::uint16_t)key.size(), key.data(),
                    0, nullptr, 0, nullptr);
        send(req.data(), req.length());
    }

private:
    const int m_socket;
    cybozu::dynbuf m_buffer;
    std::size_t m_last_response_size = 0;

    void send(const char* p, std::size_t len) {
        while( len > 0 ) {
            ssize_t n = ::send(m_socket, p, len, 0);
            cybozu_assert( n != -1 );
            p += n;
            len -= n;
        }
    }

    bool recv() {
        char* p = m_buffer.prepare(256 << 10);
        ssize_t n = ::recv(m_socket, p, 256<<10, 0);
        cybozu_assert( n != -1 );
        if( n == 0 ) return false;
        m_buffer.consume(n);
        return true;
    }
};


// tests
AUTOTEST(opaque) {
    client c;
    response r;

    c.noop(nullptr);
    cybozu_assert( c.get_response(r) );
    cybozu_assert( r.command() == binary_command::Noop );
    cybozu_assert( r.status() == binary_status::OK );
    opaque_t zero = {'\0', '\0', '\0', '\0'};
    cybozu_assert( std::memcmp(r.opaque(), &zero, sizeof(opaque_t)) == 0 );

    opaque_t op1 = {'\x12', '\x23', '\x45', '\x67'};
    c.noop(&op1);
    cybozu_assert( c.get_response(r) );
    cybozu_assert( r.command() == binary_command::Noop );
    cybozu_assert( r.status() == binary_status::OK );
    cybozu_assert( std::memcmp(r.opaque(), &op1, sizeof(opaque_t)) == 0 );
}


// main
bool optparse(int argc, char** argv) {
    if( argc != 2 && argc != 3 ) {
        std::cout << "Usage: protocol_binary SERVER [PORT]" << std::endl;
        return false;
    }
    g_server = argv[1];
    if( argc == 3 ) {
        int n = std::stoi(argv[2]);
        if( n <= 0 || n > 65535 ) {
            std::cout << "Invalid port number: " << argv[2] << std::endl;
            return false;
        }
        g_port = n;
    }

    int s = connect_server();
    if( s == -1 ) {
        std::cout << "Failed to connect to " << g_server << std::endl;
        return false;
    }
    ::close(s);
    return true;
}

TEST_MAIN(optparse);