stub_impl.hpp

#ifndef BAIDU_PADDLE_SERVING_SDK_CPP_STUB_IMPL_HPP
#define BAIDU_PADDLE_SERVING_SDK_CPP_STUB_IMPL_HPP

namespace baidu {
namespace paddle_serving {
namespace sdk_cpp {

template<typename T, typename C, typename R, typename I, typename O> 
int StubImpl<T, C, R, I, O>::initialize(
        const VariantInfo& var, const std::string& ep,
        const std::string* tag, const std::string* tag_value) {
    if (tag != NULL && tag_value != NULL) {
        TagFilter* filter = new (std::nothrow) TagFilter(
                *tag, *tag_value);
        if (!filter) {
            LOG(FATAL) << "Failed create tag filter, key: " << tag 
                << ", value: " << tag_value;
            return -1;
        }

        _gchannel = init_channel(var, filter);
        LOG(INFO) 
            << "Create stub with tag: " << *tag
            << ", " << *tag_value << ", ep: " << ep;
    } else {
        _gchannel = init_channel(var, NULL);
        LOG(INFO) << "Create stub without tag, ep " << ep;
    }

    if (!_gchannel) {
        LOG(FATAL) << "Failed init channel via var_info";
        return -1;
    }

    _service_stub = new (std::nothrow) T(_gchannel);
    if (!_service_stub) {
        LOG(FATAL) << "Failed create stub with channel";
        return -1;
    }

    _infer = _service_stub->GetDescriptor()->FindMethodByName(
            INFERENCE_METHOD_NAME);
    if (!_infer) {
        LOG(FATAL) << "Failed get inference method, "
            << "method name: " << INFERENCE_METHOD_NAME;
        return -1;
    }

    _debug = _service_stub->GetDescriptor()->FindMethodByName(
            DEBUG_METHOD_NAME);
    if (!_debug) {
        LOG(FATAL) << "Failed get debug method, "
            << "method name: " << DEBUG_METHOD_NAME;
        return -1;
    }

    _endpoint = ep;

    if (bthread_key_create(&_bthread_key, NULL) != 0) {
        LOG(FATAL) << "Failed create key for stub tls"; 
        return -1;
    }

    const std::string& name 
        = _endpoint + "_" + _service_stub->GetDescriptor()->full_name() + "_" + _tag;

    _ltc_bvars.clear();
    _avg_bvars.clear();
    BAIDU_SCOPED_LOCK(_bvar_mutex);

#ifndef DEFINE_LATENCY
#define DEFINE_LATENCY(item)                                \
    do {                                                    \
        _ltc_##item = new (std::nothrow) LatencyWrapper(name + "_"#item);\
        if (!_ltc_##item) {                                 \
            LOG(FATAL) << "Failed create latency recorder:" \
                << name + "_"#item;                         \
            return -1;                                      \
        }                                                   \
        _ltc_bvars["ltc_"#item] = _ltc_##item;              \
    } while(0)
#endif

    DEFINE_LATENCY(infer_sync);
    DEFINE_LATENCY(infer_async);
    DEFINE_LATENCY(infer_send);
    DEFINE_LATENCY(infer_recv);
    DEFINE_LATENCY(infer_cancel);
    DEFINE_LATENCY(debug);
    DEFINE_LATENCY(rpc_init);
    DEFINE_LATENCY(thrd_clear);
    DEFINE_LATENCY(pack_map);
    DEFINE_LATENCY(pack_merge);

#undef DEFINE_LATENCY

#ifndef DEFINE_AVERAGE
#define DEFINE_AVERAGE(item)                                    \
    do {                                                        \
        _avg_##item = new(std::nothrow) AverageWrapper(name + "_"#item);\
        if (!_avg_##item) {                                     \
            LOG(FATAL) << "Failed create average recorder:"     \
                << name + "_"#item;                             \
            return -1;                                          \
        }                                                       \
        _avg_bvars["avg_"#item] = _avg_##item;                  \
    } while(0)
#endif

    DEFINE_AVERAGE(failure);
    DEFINE_AVERAGE(pack);
    DEFINE_AVERAGE(item_size);
    DEFINE_AVERAGE(pack_fail);

#undef DEFINE_AVERAGE

    return 0;
}

template<typename T, typename C, typename R, typename I, typename O> 
int StubImpl<T, C, R, I, O>::thrd_initialize() {
    if (bthread_getspecific(_bthread_key) != NULL) {
        LOG(WARNING) << "Already thread initialized for stub"; 
        return 0;
    }

    StubTLS* tls = new (std::nothrow) StubTLS();
    if (!tls || bthread_setspecific(_bthread_key, tls) != 0) {
        LOG(FATAL) << "Failed binding tls data to bthread_key";
        return -1;
    }

    LOG(WARNING) << "Succ thread initialize stub impl!";

    return 0;
}

template<typename T, typename C, typename R, typename I, typename O> 
int StubImpl<T, C, R, I, O>::thrd_clear() {
    MetricScope metric(this, "thrd_clear");
    StubTLS* tls = get_tls();
    if (!tls) {
        LOG(FATAL) << "Failed get tls stub object"; 
        return -1;
    }

    // clear predictor
    size_t ps = tls->predictor_pools.size();
    for (size_t pi = 0; pi < ps; ++pi) {
        Predictor* p = tls->predictor_pools[pi];
        if (p && p->is_inited() && return_predictor(p) != 0) {
            LOG(FATAL) << "Failed return predictor: " << pi;
            return -1;
        }
    }
    tls->predictor_pools.clear();

    // clear request
    size_t is = tls->request_pools.size();
    for (size_t ii = 0; ii < is; ++ii) {
        if (return_request(tls->request_pools[ii])!= 0)  {
            LOG(FATAL) << "Failed return request: " << ii; 
            return -1;
        }
    }
    tls->request_pools.clear();

    // clear response
    size_t os = tls->response_pools.size();
    for (size_t oi = 0; oi < os; ++oi) {
        if (return_response(tls->response_pools[oi])!= 0)  {
            LOG(FATAL) << "Failed return response: " << oi;
            return -1;
        }
    }
    tls->response_pools.clear();
    return 0;
}

template<typename T, typename C, typename R, typename I, typename O> 
int StubImpl<T, C, R, I, O>::thrd_finalize() {
    StubTLS* tls = get_tls();
    if (!tls || thrd_clear() != 0) {
        LOG(FATAL) << "Failed clreate tls in thrd finalize";
        return -1;
    }
    
    delete tls; 
    return 0;
}

template<typename T, typename C, typename R, typename I, typename O> 
Predictor* StubImpl<T, C, R, I, O>::fetch_predictor() {
    StubTLS* tls = get_tls();
    if (!tls) {
        LOG(FATAL) << "Failed get tls data when fetching predictor"; 
        return NULL;
    }

    PredictorImpl<T>* predictor = butil::get_object<PredictorImpl<T> >();
    if (!predictor) {
        LOG(FATAL) << "Failed fetch predictor";
        return NULL;
    }

    if (predictor->init(_gchannel, _service_stub, _infer, _debug, _options, 
                this, _tag) != 0) {
        LOG(FATAL) << "Failed init fetched predictor"; 
        return NULL;
    }

    tls->predictor_pools.push_back(predictor);
    return predictor;
}

template<typename T, typename C, typename R, typename I, typename O> 
int StubImpl<T, C, R, I, O>::return_predictor(Predictor* predictor) {
    if ((dynamic_cast<PredictorImpl<T>*>(predictor))->deinit() != 0) {
        LOG(FATAL) << "Failed deinit fetched predictor"; 
        return -1;
    }
    butil::return_object(dynamic_cast<PredictorImpl<T>*>(predictor));
    return 0;
}

template<typename T, typename C, typename R, typename I, typename O> 
int StubImpl<T, C, R, I, O>::return_predictor(Predictor* predictor) const {
    if ((dynamic_cast<PredictorImpl<T>*>(predictor))->deinit() != 0) {
        LOG(FATAL) << "Failed deinit fetched predictor"; 
        return -1;
    }
    butil::return_object(dynamic_cast<PredictorImpl<T>*>(predictor));
    return 0;
}

template<typename T, typename C, typename R, typename I, typename O> 
google::protobuf::Message* StubImpl<T, C, R, I, O>::fetch_request() {
    StubTLS* tls = get_tls();
    if (!tls) {
        LOG(FATAL) << "Failed get tls data when fetching request"; 
        return NULL;
    }

    I* req = butil::get_object<I>();
    if (!req) {
        LOG(FATAL) << "Failed get tls request item, type: " << typeid(I).name();
        return NULL;
    }

    req->Clear();
    tls->request_pools.push_back(req);
    return req;
}

template<typename T, typename C, typename R, typename I, typename O>
int StubImpl<T, C, R, I, O>::return_request(
        google::protobuf::Message* request) const {
    request->Clear();
    butil::return_object(dynamic_cast<I*>(request));
    return 0;
}

template<typename T, typename C, typename R, typename I, typename O> 
int StubImpl<T, C, R, I, O>::return_request(
        google::protobuf::Message* request) {
    request->Clear();
    butil::return_object(dynamic_cast<I*>(request));
    return 0;
}

template<typename T, typename C, typename R, typename I, typename O> 
google::protobuf::Message* StubImpl<T, C, R, I, O>::fetch_response() {
    StubTLS* tls = get_tls();
    if (!tls) {
        LOG(FATAL) << "Failed get tls data when fetching response"; 
        return NULL;
    }

    O* res = butil::get_object<O>();
    if (!res) {
        LOG(FATAL) << "Failed get tls response item, type: " << typeid(O).name();
        return NULL;
    }

    res->Clear();
    tls->response_pools.push_back(res);
    return res;
}

template<typename T, typename C, typename R, typename I, typename O> 
int StubImpl<T, C, R, I, O>::return_response(
        google::protobuf::Message* response) const {
    response->Clear();
    butil::return_object(dynamic_cast<O*>(response));
    return 0;
}

template<typename T, typename C, typename R, typename I, typename O> 
int StubImpl<T, C, R, I, O>::return_response(
        google::protobuf::Message* response) {
    response->Clear();
    butil::return_object(dynamic_cast<O*>(response));
    return 0;
}

template<typename T, typename C, typename R, typename I, typename O> 
google::protobuf::RpcChannel* StubImpl<T, C, R, I, O>::init_channel(
        const VariantInfo& var, brpc::NamingServiceFilter* filter) {
    brpc::ChannelOptions chn_options; 
    chn_options.ns_filter = filter;

    // parameters
    ASSIGN_CONF_ITEM(chn_options.protocol, var.parameters.protocol, NULL);
    ASSIGN_CONF_ITEM(_tag, var.parameters.route_tag, NULL);
    ASSIGN_CONF_ITEM(_max_channel, var.parameters.max_channel, NULL);
    ASSIGN_CONF_ITEM(_package_size, var.parameters.package_size, NULL);

    if (_max_channel < 1) {
        LOG(ERROR) << "Invalid MaxChannelPerRequest: " << _max_channel;
        return NULL;
    }

    // connection
    ASSIGN_CONF_ITEM(chn_options.max_retry, var.connection.cnt_retry_conn, NULL);
    ASSIGN_CONF_ITEM(chn_options.connect_timeout_ms, var.connection.tmo_conn, NULL);
    ASSIGN_CONF_ITEM(chn_options.timeout_ms, var.connection.tmo_rpc, NULL);
    ASSIGN_CONF_ITEM(chn_options.backup_request_ms, var.connection.tmo_hedge, NULL);

    // connection type
    std::string conn_type_str;
    ASSIGN_CONF_ITEM(conn_type_str, var.connection.type_conn, NULL);
    chn_options.connection_type 
        = brpc::StringToConnectionType(conn_type_str);

    // naminginfo
    std::string cluster_naming_info;
    std::string cluster_loadbalancer;
    ASSIGN_CONF_ITEM(cluster_naming_info, var.naminginfo.cluster_naming, NULL);
    ASSIGN_CONF_ITEM(cluster_loadbalancer, var.naminginfo.load_balancer, NULL);

    // brpc single channel
    _channel = butil::get_object<brpc::Channel>();
    if (!_channel) {
        LOG(FATAL) << "Failed get channel object from butil::pool";
        return NULL;
    }

    if (_channel->Init(
                cluster_naming_info.c_str(), 
                cluster_loadbalancer.c_str(), 
                &chn_options) != 0) {
        LOG(ERROR) 
            << "Failed to initialize channel, path: " 
            << cluster_naming_info;
        return NULL;
    }

    // brpc parallel channel
    _pchannel = init_pchannel(
            _channel, _max_channel, _package_size, chn_options);
    if (_pchannel) {
        LOG(INFO) << "Succ create parallel channel, count: "
            << _max_channel;
        return _pchannel;
    }

    return _channel;
}

template<typename T, typename C, typename R, typename I, typename O> 
brpc::ParallelChannel* StubImpl<T, C, R, I, O>::init_pchannel(
        brpc::Channel* sub_channel, uint32_t channel_count,
        uint32_t package_size, const brpc::ChannelOptions& options) {
    if (channel_count <= 1) { // noneed use parallel channel
        LOG(INFO) << "channel count <= 1, noneed use pchannel.";
        return NULL;
    }

    _pchannel = butil::get_object<brpc::ParallelChannel>();
    if (!_pchannel) {
        LOG(FATAL) << "Failed get pchannel from object pool";
        return NULL;
    }

    brpc::ParallelChannelOptions pchan_options;
    pchan_options.timeout_ms = options.timeout_ms;
    if (_pchannel->Init(&pchan_options) != 0) {
        LOG(FATAL) << "Failed init parallel channel with tmo_us: "
            << pchan_options.timeout_ms;
        return NULL;
    }

    for (uint32_t si = 0; si < channel_count; ++si) {
        if (_pchannel->AddChannel(
                    sub_channel, 
                    brpc::DOESNT_OWN_CHANNEL,
                    new C(package_size, this), 
                    new R(package_size, this)) != 0) {
            LOG(FATAL) << "Failed add channel at: " << si
                << ", package_size:" << package_size;
            return NULL;
        }
    }

    return _pchannel;
}

} // sdk_cpp
} // paddle_serving
} // baidu

#endif  //BAIDU_PADDLE_SERVING_SDK_CPP_STUB_IMPL_HPP