// Baidu RPC - A framework to host and access services throughout Baidu.
// Copyright (c) 2015 Baidu.com, Inc. All Rights Reserved

// Author: The baidu-rpc authors (pbrpc@baidu.com)
// Date: Fri Jun  5 18:25:40 CST 2015

#include <google/protobuf/descriptor.h>         // MethodDescriptor
#include <google/protobuf/message.h>            // Message
#include <gflags/gflags.h>
#include "base/logging.h"                       // LOG()
#include "base/time.h"
#include "base/iobuf.h"                         // base::IOBuf
#include "brpc/controller.h"               // Controller
#include "brpc/details/controller_private_accessor.h"
#include "brpc/socket.h"                   // Socket
#include "brpc/server.h"                   // Server
#include "brpc/details/server_private_accessor.h"
#include "brpc/span.h"
#include "brpc/redis.h"
#include "brpc/policy/redis_protocol.h"


namespace brpc {

DECLARE_bool(enable_rpcz);

namespace policy {

DEFINE_bool(redis_verbose, false,
            "[DEBUG] Print EVERY redis request/response to stderr");

struct InputResponse : public InputMessageBase {
    bthread_id_t id_wait;
    RedisResponse response;

    // @InputMessageBase
    void DestroyImpl() {
        delete this;
    }
};

// "Message" = "Response" as we only implement the client for redis.
ParseResult ParseRedisMessage(base::IOBuf* source, Socket* socket,
                              bool /*read_eof*/, const void* /*arg*/) {
    if (source->empty()) {
        return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
    }
    // NOTE(gejun): PopPipelinedInfo() is actually more contended than what
    // I thought before. The Socket._pipeline_q is a SPSC queue pushed before
    // sending and popped when response comes back, being protected by a
    // mutex. Previously the mutex is shared with Socket._id_wait_list. When
    // 200 bthreads access one redis-server, ~1.5s in total is spent on
    // contention in 10-second duration. If the mutex is separated, the time
    // drops to ~0.25s. I further replaced PeekPipelinedInfo() with
    // GivebackPipelinedInfo() to lock only once(when receiving response)
    // in most cases, and the time decreases to ~0.14s.
    PipelinedInfo pi;
    if (!socket->PopPipelinedInfo(&pi)) {
        LOG(WARNING) << "No corresponding PipelinedInfo in socket";
        return MakeParseError(PARSE_ERROR_TRY_OTHERS);
    }
    InputResponse* msg = static_cast<InputResponse*>(socket->parsing_context());
    if (msg == NULL) {
        msg = new InputResponse;
        socket->reset_parsing_context(msg);
    }

    if (!msg->response.ConsumePartialIOBuf(*source, pi.count)) {
        socket->GivebackPipelinedInfo(pi);
        return MakeParseError(PARSE_ERROR_NOT_ENOUGH_DATA);
    }
    CHECK_EQ((uint32_t)msg->response.reply_size(), pi.count);
    msg->id_wait = pi.id_wait;
    socket->release_parsing_context();
    return MakeMessage(msg);
}

void ProcessRedisResponse(InputMessageBase* msg_base) {
    const int64_t start_parse_us = base::cpuwide_time_us();
    DestroyingPtr<InputResponse> msg(static_cast<InputResponse*>(msg_base));

    const bthread_id_t cid = msg->id_wait;
    Controller* cntl = NULL;
    const int rc = bthread_id_lock(cid, (void**)&cntl);
    if (rc != 0) {
        LOG_IF(FATAL, rc != EINVAL) << "Fail to lock correlation_id="
                                    << cid.value << ": " << berror(rc);
        return;
    }
    
    ControllerPrivateAccessor accessor(cntl);
    Span* span = accessor.span();
    if (span) {
        span->set_base_real_us(msg->base_real_us());
        span->set_received_us(msg->received_us());
        span->set_response_size(msg->response.ByteSize());
        span->set_start_parse_us(start_parse_us);
    }
    const int saved_error = cntl->ErrorCode();
    if (cntl->response() != NULL) {
        if (cntl->response()->GetDescriptor() != RedisResponse::descriptor()) {
            cntl->SetFailed(ERESPONSE, "Must be RedisResponse");
        } else {
            // We work around ParseFrom of pb which is just a placeholder.
            if (msg->response.reply_size() != (int)accessor.pipelined_count()) {
                cntl->SetFailed(ERESPONSE, "pipelined_count=%d of response does "
                                      "not equal request's=%d",
                                      msg->response.reply_size(), accessor.pipelined_count());
            }
            ((RedisResponse*)cntl->response())->Swap(&msg->response);
            if (FLAGS_redis_verbose) {
                std::cerr << "[REDIS RESPONSE] "
                          << *((RedisResponse*)cntl->response())
                          << std::endl;
            }
        }
    } // silently ignore the response.
    
    // Unlocks correlation_id inside. Revert controller's
    // error code if it version check of `cid' fails
    msg.reset();  // optional, just release resourse ASAP
    accessor.OnResponse(cid, saved_error);
}

void SerializeRedisRequest(base::IOBuf* buf,
                           Controller* cntl,
                           const google::protobuf::Message* request) {
    if (request == NULL) {
        return cntl->SetFailed(EREQUEST, "request is NULL");
    }
    if (request->GetDescriptor() != RedisRequest::descriptor()) {
        return cntl->SetFailed(EREQUEST, "The request is not a RedisRequest");
    }
    const RedisRequest* rr = (const RedisRequest*)request;
    // We work around SerializeTo of pb which is just a placeholder.
    if (!rr->SerializeTo(buf)) {
        return cntl->SetFailed(EREQUEST, "Fail to serialize RedisRequest");
    }
    ControllerPrivateAccessor(cntl).set_pipelined_count(rr->command_size());
    if (FLAGS_redis_verbose) {
        std::cerr << "[REDIS REQUEST] " << *rr << std::endl;
    }
}

void PackRedisRequest(base::IOBuf* buf,
                      SocketMessage**,
                      uint64_t /*correlation_id*/,
                      const google::protobuf::MethodDescriptor*,
                      Controller*,
                      const base::IOBuf& request,
                      const Authenticator* /*auth*/) {
    buf->append(request);
}

const std::string& GetRedisMethodName(
    const google::protobuf::MethodDescriptor*,
    const Controller*) {
    const static std::string REDIS_SERVER_STR = "redis-server";
    return REDIS_SERVER_STR;
}

}  // namespace policy
} // namespace brpc

