ascii_scrub.c

/* -*- Mode: C; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
/*
 * arcus-memcached - Arcus memory cache server
 * Copyright 2010-2014 NAVER Corp.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "protocol_extension.h"

static const char *get_name(const void *cmd_cookie);
static bool accept_command(const void *cmd_cookie, void *cookie,
                           int argc, token_t *argv, size_t *ndata,
                           char **ptr);
static bool execute_command(const void *cmd_cookie, const void *cookie,
                            int argc, token_t *argv,
                            bool (*response_handler)(const void *cookie,
                                                     int nbytes,
                                                     const char *dta));
static void abort_command(const void *cmd_cookie, const void *cookie);

static EXTENSION_ASCII_PROTOCOL_DESCRIPTOR scrub_descriptor = {
    .get_name = get_name,
    .accept = accept_command,
    .execute = execute_command,
    .abort = abort_command,
    .cookie = &scrub_descriptor
};

GET_SERVER_API server_api;

static const char *get_name(const void *cmd_cookie) {
    return "scrub";
}

static bool accept_command(const void *cmd_cookie, void *cookie,
                           int argc, token_t *argv, size_t *ndata,
                           char **ptr) {
    if (strcmp(argv[0].value, "scrub") == 0) {
        if (argc == 1) return true;
        if (argc == 2 && strcmp(argv[1].value, "stale") == 0) return true;
    }
    return false;
}

static bool my_response_handler(const void *key, uint16_t keylen,
                                const void *ext, uint8_t extlen,
                                const void *body, uint32_t bodylen,
                                uint8_t datatype, uint16_t status,
                                uint64_t cas, const void *cookie)
{
    uint16_t *rval = (uint16_t*)cookie;
    *rval = status;
    return true;
}

static bool execute_command(const void *cmd_cookie, const void *cookie,
                            int argc, token_t *argv,
                            bool (*response_handler)(const void *cookie,
                                                     int nbytes,
                                                     const char *dta)) {

    protocol_binary_request_header request = {
        .request.magic = (uint8_t)PROTOCOL_BINARY_REQ,
        .request.opcode = PROTOCOL_BINARY_CMD_SCRUB
    };

    uint16_t status = 0;
    SERVER_HANDLE_V1 *server = server_api();
    ENGINE_HANDLE_V1 *v1 = (ENGINE_HANDLE_V1*)server->engine;
    if (v1 == NULL) {
        return response_handler(cookie, 29, "SERVER_ERROR internal error\r\n");
    }

    if (argc == 2 && strcmp(argv[1].value, "stale") == 0) {
        request.request.opcode = PROTOCOL_BINARY_CMD_SCRUB_STALE;
    }

    v1->unknown_command(server->engine, &status, &request, my_response_handler);

    if (status == PROTOCOL_BINARY_RESPONSE_SUCCESS) {
        return response_handler(cookie, 4, "OK\r\n");
    } else if (status == PROTOCOL_BINARY_RESPONSE_EBUSY) {
        return response_handler(cookie, 6, "BUSY\r\n");
    } else {
        return response_handler(cookie, 15, "NOT_SUPPORTED\r\n");
    }
}

static void abort_command(const void *cmd_cookie, const void *cookie)
{
    /* EMPTY */
}

MEMCACHED_PUBLIC_API
EXTENSION_ERROR_CODE memcached_extensions_initialize(const char *config,
                                                     GET_SERVER_API get_server_api) {
    server_api = get_server_api;
    SERVER_HANDLE_V1 *server = get_server_api();
    if (server == NULL) {
        return EXTENSION_FATAL;
    }
    if (!server->extension->register_extension(EXTENSION_ASCII_PROTOCOL,
                                               &scrub_descriptor)) {
        return EXTENSION_FATAL;
    }

    return EXTENSION_SUCCESS;
}