<commit_before>/*
 * memcached client implementation.
 *
 * author: <NAME> <<EMAIL>>
 */

#include "memcached-client.h"
#include "memcached-packet.h"
#include "lease.h"
#include "async.h"
#include "event2.h"
#include "istream-internal.h"
#include "buffered-io.h"

#include <daemon/log.h>

#include <glib.h>

#include <errno.h>
#include <sys/socket.h>
#include <string.h>

struct memcached_client {
    pool_t pool;

    /* I/O */
    int fd;
    struct lease_ref lease_ref;
    struct event2 event;

    /* request */
    struct {
        memcached_response_handler_t handler;
        void *handler_ctx;

        struct async_operation async;

        istream_t istream;
    } request;

    /* response */
    struct {
        enum {
            READ_HEADER,
            READ_EXTRAS,
            READ_KEY,
            READ_VALUE,
            READ_END,
        } read_state;

        struct memcached_response_header header;

        struct fifo_buffer *input;

        unsigned char *extras;

        struct {
            void *buffer;
            unsigned char *tail;
            size_t remaining;
        } key;

        struct istream value;
        size_t value_remaining;
    } response;
};

static const struct timeval timeout = {
    .tv_sec = 5,
    .tv_usec = 0,
};

static inline bool
memcached_connection_valid(const struct memcached_client *client)
{
    return client->fd >= 0;
}

/**
 * Release resources held by this object: the event object, the socket
 * lease, and the pool reference.
 */
static void
memcached_client_release(struct memcached_client *client, bool reuse)
{
    assert(client != NULL);

    event2_set(&client->event, 0);
    event2_commit(&client->event);
    client->fd = -1;
    lease_release(&client->lease_ref, reuse);
    pool_unref(client->pool);
}

static void
memcached_connection_close(struct memcached_client *client)
{
    if (!memcached_connection_valid(client))
        return;

    pool_ref(client->pool);

    event2_set(&client->event, 0);
    event2_commit(&client->event);

    switch (client->response.read_state) {
    case READ_HEADER:
    case READ_EXTRAS:
    case READ_KEY:
        client->request.handler(-1, NULL, 0, NULL, 0, NULL,
                                client->request.handler_ctx);
        client->response.read_state = READ_END;
        break;

    case READ_VALUE:
        istream_deinit_abort(&client->response.value);
        client->response.read_state = READ_END;
        break;

    case READ_END:
        break;
    }

    if (client->request.istream != NULL)
        istream_free_handler(&client->request.istream);

    if (client->fd >= 0)
        memcached_client_release(client, false);

    pool_unref(client->pool);
}

static bool
memcached_consume_input(struct memcached_client *client);

static void
memcached_client_try_read(struct memcached_client *client);

/*
 * response value istream
 *
 */

static inline struct memcached_client *
istream_to_memcached_client(istream_t istream)
{
    return (struct memcached_client *)(((char*)istream) - offsetof(struct memcached_client, response.value));
}

static off_t
istream_memcached_available(istream_t istream, G_GNUC_UNUSED bool partial)
{
    struct memcached_client *client = istream_to_memcached_client(istream);

    assert(client->response.read_state == READ_VALUE);

    return client->response.value_remaining;
}

static void
istream_memcached_read(istream_t istream)
{
    struct memcached_client *client = istream_to_memcached_client(istream);

    assert(client->response.read_state == READ_VALUE);

    if (!fifo_buffer_empty(client->response.input))
        memcached_consume_input(client);
    else {
        pool_ref(client->pool);
        memcached_client_try_read(client);
        pool_unref(client->pool);
    }
}

static void
istream_memcached_close(istream_t istream)
{
    struct memcached_client *client = istream_to_memcached_client(istream);

    assert(client->response.read_state == READ_VALUE);

    istream_deinit_abort(&client->response.value);
    memcached_client_release(client, false);
}

static const struct istream memcached_response_value = {
    .available = istream_memcached_available,
    .read = istream_memcached_read,
    .close = istream_memcached_close,
};

/*
 * response parser
 *
 */

static bool
memcached_consume_header(struct memcached_client *client)
{
    size_t length;
    const void *data = fifo_buffer_read(client->response.input, &length);

    assert(client->response.read_state == READ_HEADER);

    if (data == NULL || length < sizeof(client->response.header))
        /* not enough data yet */
        return false;

    memcpy(&client->response.header, data, sizeof(client->response.header));
    fifo_buffer_consume(client->response.input,
                        sizeof(client->response.header));
    client->response.read_state = READ_EXTRAS;

    client->response.value_remaining =
        g_ntohl(client->response.header.body_length) -
        g_ntohs(client->response.header.key_length) -
        client->response.header.extras_length;
    if (client->response.header.magic != MEMCACHED_MAGIC_RESPONSE ||
        client->response.value_remaining > G_MAXINT) {
        /* integer underflow */
        memcached_connection_close(client);
        return false;
    }

    return true;
}

static bool
memcached_consume_extras(struct memcached_client *client)
{
    assert(client->response.read_state == READ_EXTRAS);

    if (client->response.header.extras_length > 0) {
        size_t length;
        const void *data = fifo_buffer_read(client->response.input, &length);

        if (data == NULL ||
            length < sizeof(client->response.header.extras_length))
            return false;

        client->response.extras = p_malloc(client->pool,
                                           client->response.header.extras_length);
        memcpy(client->response.extras, data,
               client->response.header.extras_length);

        fifo_buffer_consume(client->response.input,
                            client->response.header.extras_length);
    } else
        client->response.extras = NULL;

    client->response.read_state = READ_KEY;
    client->response.key.remaining =
        g_ntohs(client->response.header.key_length);
    if (client->response.key.remaining > 0)
        client->response.key.buffer
            = client->response.key.tail
            = p_malloc(client->pool, client->response.key.remaining);
    else
        client->response.key.buffer = NULL;

    return true;
}

static bool
memcached_consume_key(struct memcached_client *client)
{
    assert(client->response.read_state == READ_KEY);

    if (client->response.key.remaining > 0) {
        size_t length;
        const void *data = fifo_buffer_read(client->response.input, &length);

        if (data == NULL)
            return false;

        if (length > client->response.key.remaining)
            length = client->response.key.remaining;

        memcpy(client->response.key.tail, data, length);
        client->response.key.tail += length;
        client->response.key.remaining -= length;

        if (client->response.key.remaining > 0)
            return false;
    }

    if (client->response.value_remaining > 0) {
        /* there's a value: pass it to the callback, continue
           reading */
        istream_t value;
        bool valid;

        client->response.read_state = READ_VALUE;

        istream_init(&client->response.value, &memcached_response_value,
                     client->pool);
        value = istream_struct_cast(&client->response.value);

        pool_ref(client->pool);
        client->request.handler(g_ntohs(client->response.header.status),
                                client->response.extras,
                                client->response.header.extras_length,
                                client->response.key.buffer,
                                g_ntohs(client->response.header.key_length),
                                value, client->request.handler_ctx);

        /* check if the callback has closed the value istream */
        valid = memcached_connection_valid(client);
        pool_unref(client->pool);

        return valid;
    } else {
        /* no value: invoke the callback, quit */
        client->response.read_state = READ_END;

        client->request.handler(g_ntohs(client->response.header.status),
                                client->response.extras,
                                client->response.header.extras_length,
                                client->response.key.buffer,
                                g_ntohs(client->response.header.key_length),
                                NULL, client->request.handler_ctx);

        memcached_connection_close(client);
        return false;
    }
}

static bool
memcached_consume_value(struct memcached_client *client)
{
    size_t length, nbytes;
    const void *data;

    assert(client->response.read_state == READ_VALUE);
    assert(client->response.value_remaining > 0);

    data = fifo_buffer_read(client->response.input, &length);
    if (data == NULL)
        return false;

    if (length > client->response.value_remaining)
        length = client->response.value_remaining;

    nbytes = istream_invoke_data(&client->response.value, data, length);
    if (nbytes == 0)
        return false;

    fifo_buffer_consume(client->response.input, nbytes);

    client->response.value_remaining -= nbytes;
    if (client->response.value_remaining > 0)
        return false;

    client->response.read_state = READ_END;
    istream_deinit_eof(&client->response.value);
    memcached_connection_close(client);
    return true;
}

static bool
memcached_consume_input(struct memcached_client *client)
{
    if (client->response.read_state == READ_HEADER &&
        !memcached_consume_header(client))
        return false;

    if (client->response.read_state == READ_EXTRAS &&
        !memcached_consume_extras(client))
        return false;

    if (client->response.read_state == READ_KEY &&
        !memcached_consume_key(client))
        return false;

    assert(client->response.read_state == READ_VALUE);

    return memcached_consume_value(client);
}

static void
memcached_client_try_read(struct memcached_client *client)
{
    ssize_t nbytes;

    assert(memcached_connection_valid(client));

    nbytes = recv_to_buffer(client->fd, client->response.input,
                            G_MAXINT);
    assert(nbytes != -2);

    if (nbytes == 0) {
        daemon_log(1, "memcached server closed the connection\n");
        memcached_connection_close(client);
        return;
    }

    if (nbytes < 0) {
        if (errno == EAGAIN) {
            event2_or(&client->event, EV_READ);
            return;
        }

        daemon_log(1, "read error on memcached connection: %s\n",
                   strerror(errno));
        memcached_connection_close(client);
        return;
    }

    memcached_consume_input(client);

    if (memcached_connection_valid(client) &&
        !fifo_buffer_full(client->response.input))
        event2_setbit(&client->event, EV_READ,
                      !fifo_buffer_full(client->response.input));
}

/**
 * The libevent callback for the socket.
 */
static void
memcached_client_event_callback(G_GNUC_UNUSED int fd, short event, void *ctx)
{
    struct memcached_client *client = ctx;

    if (unlikely(event & EV_TIMEOUT)) {
        daemon_log(4, "timeout\n");
        memcached_connection_close(client);
        return;
    }

    pool_ref(client->pool);
<commit_msg>memcached-client: simplified EV_READ setting in try_read()<commit_after>/*
 * memcached client implementation.
 *
 * author: <NAME> <mk<EMAIL>>
 */

#include "memcached-client.h"
#include "memcached-packet.h"
#include "lease.h"
#include "async.h"
#include "event2.h"
#include "istream-internal.h"
#include "buffered-io.h"

#include <daemon/log.h>

#include <glib.h>

#include <errno.h>
#include <sys/socket.h>
#include <string.h>

struct memcached_client {
    pool_t pool;

    /* I/O */
    int fd;
    struct lease_ref lease_ref;
    struct event2 event;

    /* request */
    struct {
        memcached_response_handler_t handler;
        void *handler_ctx;

        struct async_operation async;

        istream_t istream;
    } request;

    /* response */
    struct {
        enum {
            READ_HEADER,
            READ_EXTRAS,
            READ_KEY,
            READ_VALUE,
            READ_END,
        } read_state;

        struct memcached_response_header header;

        struct fifo_buffer *input;

        unsigned char *extras;

        struct {
            void *buffer;
            unsigned char *tail;
            size_t remaining;
        } key;

        struct istream value;
        size_t value_remaining;
    } response;
};

static const struct timeval timeout = {
    .tv_sec = 5,
    .tv_usec = 0,
};

static inline bool
memcached_connection_valid(const struct memcached_client *client)
{
    return client->fd >= 0;
}

/**
 * Release resources held by this object: the event object, the socket
 * lease, and the pool reference.
 */
static void
memcached_client_release(struct memcached_client *client, bool reuse)
{
    assert(client != NULL);

    event2_set(&client->event, 0);
    event2_commit(&client->event);
    client->fd = -1;
    lease_release(&client->lease_ref, reuse);
    pool_unref(client->pool);
}

static void
memcached_connection_close(struct memcached_client *client)
{
    if (!memcached_connection_valid(client))
        return;

    pool_ref(client->pool);

    event2_set(&client->event, 0);
    event2_commit(&client->event);

    switch (client->response.read_state) {
    case READ_HEADER:
    case READ_EXTRAS:
    case READ_KEY:
        client->request.handler(-1, NULL, 0, NULL, 0, NULL,
                                client->request.handler_ctx);
        client->response.read_state = READ_END;
        break;

    case READ_VALUE:
        istream_deinit_abort(&client->response.value);
        client->response.read_state = READ_END;
        break;

    case READ_END:
        break;
    }

    if (client->request.istream != NULL)
        istream_free_handler(&client->request.istream);

    if (client->fd >= 0)
        memcached_client_release(client, false);

    pool_unref(client->pool);
}

static bool
memcached_consume_input(struct memcached_client *client);

static void
memcached_client_try_read(struct memcached_client *client);

/*
 * response value istream
 *
 */

static inline struct memcached_client *
istream_to_memcached_client(istream_t istream)
{
    return (struct memcached_client *)(((char*)istream) - offsetof(struct memcached_client, response.value));
}

static off_t
istream_memcached_available(istream_t istream, G_GNUC_UNUSED bool partial)
{
    struct memcached_client *client = istream_to_memcached_client(istream);

    assert(client->response.read_state == READ_VALUE);

    return client->response.value_remaining;
}

static void
istream_memcached_read(istream_t istream)
{
    struct memcached_client *client = istream_to_memcached_client(istream);

    assert(client->response.read_state == READ_VALUE);

    if (!fifo_buffer_empty(client->response.input))
        memcached_consume_input(client);
    else {
        pool_ref(client->pool);
        memcached_client_try_read(client);
        pool_unref(client->pool);
    }
}

static void
istream_memcached_close(istream_t istream)
{
    struct memcached_client *client = istream_to_memcached_client(istream);

    assert(client->response.read_state == READ_VALUE);

    istream_deinit_abort(&client->response.value);
    memcached_client_release(client, false);
}

static const struct istream memcached_response_value = {
    .available = istream_memcached_available,
    .read = istream_memcached_read,
    .close = istream_memcached_close,
};

/*
 * response parser
 *
 */

static bool
memcached_consume_header(struct memcached_client *client)
{
    size_t length;
    const void *data = fifo_buffer_read(client->response.input, &length);

    assert(client->response.read_state == READ_HEADER);

    if (data == NULL || length < sizeof(client->response.header))
        /* not enough data yet */
        return false;

    memcpy(&client->response.header, data, sizeof(client->response.header));
    fifo_buffer_consume(client->response.input,
                        sizeof(client->response.header));
    client->response.read_state = READ_EXTRAS;

    client->response.value_remaining =
        g_ntohl(client->response.header.body_length) -
        g_ntohs(client->response.header.key_length) -
        client->response.header.extras_length;
    if (client->response.header.magic != MEMCACHED_MAGIC_RESPONSE ||
        client->response.value_remaining > G_MAXINT) {
        /* integer underflow */
        memcached_connection_close(client);
        return false;
    }

    return true;
}

static bool
memcached_consume_extras(struct memcached_client *client)
{
    assert(client->response.read_state == READ_EXTRAS);

    if (client->response.header.extras_length > 0) {
        size_t length;
        const void *data = fifo_buffer_read(client->response.input, &length);

        if (data == NULL ||
            length < sizeof(client->response.header.extras_length))
            return false;

        client->response.extras = p_malloc(client->pool,
                                           client->response.header.extras_length);
        memcpy(client->response.extras, data,
               client->response.header.extras_length);

        fifo_buffer_consume(client->response.input,
                            client->response.header.extras_length);
    } else
        client->response.extras = NULL;

    client->response.read_state = READ_KEY;
    client->response.key.remaining =
        g_ntohs(client->response.header.key_length);
    if (client->response.key.remaining > 0)
        client->response.key.buffer
            = client->response.key.tail
            = p_malloc(client->pool, client->response.key.remaining);
    else
        client->response.key.buffer = NULL;

    return true;
}

static bool
memcached_consume_key(struct memcached_client *client)
{
    assert(client->response.read_state == READ_KEY);

    if (client->response.key.remaining > 0) {
        size_t length;
        const void *data = fifo_buffer_read(client->response.input, &length);

        if (data == NULL)
            return false;

        if (length > client->response.key.remaining)
            length = client->response.key.remaining;

        memcpy(client->response.key.tail, data, length);
        client->response.key.tail += length;
        client->response.key.remaining -= length;

        if (client->response.key.remaining > 0)
            return false;
    }

    if (client->response.value_remaining > 0) {
        /* there's a value: pass it to the callback, continue
           reading */
        istream_t value;
        bool valid;

        client->response.read_state = READ_VALUE;

        istream_init(&client->response.value, &memcached_response_value,
                     client->pool);
        value = istream_struct_cast(&client->response.value);

        pool_ref(client->pool);
        client->request.handler(g_ntohs(client->response.header.status),
                                client->response.extras,
                                client->response.header.extras_length,
                                client->response.key.buffer,
                                g_ntohs(client->response.header.key_length),
                                value, client->request.handler_ctx);

        /* check if the callback has closed the value istream */
        valid = memcached_connection_valid(client);
        pool_unref(client->pool);

        return valid;
    } else {
        /* no value: invoke the callback, quit */
        client->response.read_state = READ_END;

        client->request.handler(g_ntohs(client->response.header.status),
                                client->response.extras,
                                client->response.header.extras_length,
                                client->response.key.buffer,
                                g_ntohs(client->response.header.key_length),
                                NULL, client->request.handler_ctx);

        memcached_connection_close(client);
        return false;
    }
}

static bool
memcached_consume_value(struct memcached_client *client)
{
    size_t length, nbytes;
    const void *data;

    assert(client->response.read_state == READ_VALUE);
    assert(client->response.value_remaining > 0);

    data = fifo_buffer_read(client->response.input, &length);
    if (data == NULL)
        return false;

    if (length > client->response.value_remaining)
        length = client->response.value_remaining;

    nbytes = istream_invoke_data(&client->response.value, data, length);
    if (nbytes == 0)
        return false;

    fifo_buffer_consume(client->response.input, nbytes);

    client->response.value_remaining -= nbytes;
    if (client->response.value_remaining > 0)
        return false;

    client->response.read_state = READ_END;
    istream_deinit_eof(&client->response.value);
    memcached_connection_close(client);
    return true;
}

static bool
memcached_consume_input(struct memcached_client *client)
{
    if (client->response.read_state == READ_HEADER &&
        !memcached_consume_header(client))
        return false;

    if (client->response.read_state == READ_EXTRAS &&
        !memcached_consume_extras(client))
        return false;

    if (client->response.read_state == READ_KEY &&
        !memcached_consume_key(client))
        return false;

    assert(client->response.read_state == READ_VALUE);

    return memcached_consume_value(client);
}

static void
memcached_client_try_read(struct memcached_client *client)
{
    ssize_t nbytes;

    assert(memcached_connection_valid(client));

    nbytes = recv_to_buffer(client->fd, client->response.input,
                            G_MAXINT);
    assert(nbytes != -2);

    if (nbytes == 0) {
        daemon_log(1, "memcached server closed the connection\n");
        memcached_connection_close(client);
        return;
    }

    if (nbytes < 0) {
        if (errno == EAGAIN) {
            event2_or(&client->event, EV_READ);
            return;
        }

        daemon_log(1, "read error on memcached connection: %s\n",
                   strerror(errno));
        memcached_connection_close(client);
        return;
    }

    memcached_consume_input(client);

    if (memcached_connection_valid(client))
        event2_setbit(&client->event, EV_READ,
                      !fifo_buffer_full(client->response.input));
}

/**
 * The libevent callback for the socket.
 */
static void
memcached_client_event_callback(G_GNUC_UNUSED int fd, short event, void *ctx)
{
    struct memcached_client *client = ctx;

    if (unlikely(event & EV_TIMEOUT)) {
        daemon_log(4, "timeout\n");
        memcached_connection_close(client);
        return;
    }

    pool_ref(client->pool);

{'fraction_non_alphanumeric': 0.10810810810810811, 'fraction_numerical': 0.002772002772002772, 'mean_word_length': 2.619532842138266, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 272, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '80110', 'n_tokens_mistral': 7013, 'n_tokens_neox': 6809, 'n_words': 1272}