/src/brpc/details/controller_private_accessor.h

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

#ifndef BRPC_CONTROLLER_PRIVATE_ACCESSOR_H
#define BRPC_CONTROLLER_PRIVATE_ACCESSOR_H

// This is an rpc-internal file.

#include "brpc/socket.h"
#include "brpc/controller.h"
#include "brpc/stream.h"

namespace google {
namespace protobuf {
class Message;
}
}

namespace brpc {

class Span;

class AuthContext;

// A wrapper to access some private methods/fields of `Controller'
// This is supposed to be used by internal RPC protocols ONLY
class ControllerPrivateAccessor {
public:
    explicit ControllerPrivateAccessor(Controller* cntl) {
        _cntl = cntl;
    }

    void OnResponse(CallId id, int saved_error) {
        const Controller::CompletionInfo info = { id, true };
        _cntl->OnVersionedRPCReturned(info, false, saved_error);
    }

    ControllerPrivateAccessor &set_peer_id(SocketId peer_id) {
        _cntl->_current_call.peer_id = peer_id;
        return *this;
    }

    Socket* get_sending_socket() {
        return _cntl->_current_call.sending_sock.get();
    }

    int64_t real_timeout_ms() {
        return _cntl->_real_timeout_ms;
    }

    void move_in_server_receiving_sock(SocketUniquePtr& ptr) {
        CHECK(_cntl->_current_call.sending_sock == NULL);
        _cntl->_current_call.sending_sock.reset(ptr.release());
    }

    StreamUserData* get_stream_user_data() {
        return _cntl->_current_call.stream_user_data;
    }

    ControllerPrivateAccessor &set_security_mode(bool security_mode) {
        _cntl->set_flag(Controller::FLAGS_SECURITY_MODE, security_mode);
        return *this;
    }

    ControllerPrivateAccessor &set_remote_side(const butil::EndPoint& pt) {
        _cntl->_remote_side = pt;
        return *this;
    }

    ControllerPrivateAccessor &set_local_side(const butil::EndPoint& pt) {
        _cntl->_local_side = pt;
        return *this;
    }
 
    ControllerPrivateAccessor &set_auth_context(const AuthContext* ctx) {
        _cntl->set_auth_context(ctx);
        return *this;
    }

    // Overloaded set_span methods to support both shared_ptr and raw pointer
    ControllerPrivateAccessor &set_span(const std::shared_ptr<Span>& span);
    ControllerPrivateAccessor &set_span(Span* span);
    
    ControllerPrivateAccessor &set_request_protocol(ProtocolType protocol) {
        _cntl->_request_protocol = protocol;
        return *this;
    }
    
    std::shared_ptr<Span> span() const;

    uint32_t pipelined_count() const { return _cntl->_pipelined_count; }
    void set_pipelined_count(uint32_t count) {  _cntl->_pipelined_count = count; }

    // The mysql protocol stores its statement type (MYSQL_NORMAL_STATEMENT /
    // MYSQL_PREPARED_STATEMENT) in the pipelined_count slot.
    void set_mysql_statement_type(uint32_t type) { set_pipelined_count(type); }

    ControllerPrivateAccessor& set_server(const Server* server) {
        _cntl->_server = server;
        return *this;
    }

    // Pass the owership of |settings| to _cntl, while is going to be
    // destroyed in Controller::Reset()
    void set_remote_stream_settings(StreamSettings *settings) {
        _cntl->_remote_stream_settings = settings;
    }
    StreamSettings* remote_stream_settings() {
        return _cntl->_remote_stream_settings;
    }

    StreamIds request_streams() { return _cntl->_request_streams; }
    StreamIds response_streams() { return _cntl->_response_streams; }

    void set_method(const google::protobuf::MethodDescriptor* method) 
    { _cntl->_method = method; }

    void set_readable_progressive_attachment(ReadableProgressiveAttachment* s)
    { _cntl->_rpa.reset(s); }

    void set_auth_flags(uint32_t auth_flags) {
        _cntl->_auth_flags = auth_flags;
    }

    void clear_auth_flags() { _cntl->_auth_flags = 0; }

    // Set how the sending socket is reserved after the RPC (mysql transactions).
    void set_bind_sock_action(BindSockAction action) { _cntl->set_bind_sock_action(action); }
    // Transfer ownership of the reserved socket to `ptr`.
    void get_bind_sock(SocketUniquePtr* ptr) {
        if (_cntl->_bind_sock) {
            _cntl->_bind_sock->ReAddress(ptr);
        }
    }
    // Reuse an externally-reserved socket for the next RPC.
    void use_bind_sock(SocketId sock_id) {
        _cntl->set_bind_sock_action(BIND_SOCK_USE);
        Socket::Address(sock_id, &_cntl->_bind_sock);
    }
    void set_session_data(void* d) { _cntl->_session_data = d; }
    void* session_data() const { return _cntl->_session_data; }

    std::string& protocol_param() { return _cntl->protocol_param(); }
    const std::string& protocol_param() const { return _cntl->protocol_param(); }

    // Note: This function can only be called in server side. The deadline of client
    // side is properly set in the RPC sending path.
    void set_deadline_us(int64_t deadline_us) { _cntl->_deadline_us = deadline_us; }

    ControllerPrivateAccessor& set_begin_time_us(int64_t begin_time_us) {
        _cntl->_begin_time_us = begin_time_us;
        _cntl->_end_time_us = UNSET_MAGIC_NUM;
        return *this;
    }

    ControllerPrivateAccessor& set_health_check_call() {
        _cntl->add_flag(Controller::FLAGS_HEALTH_CHECK_CALL);
        return *this;
    }

    void set_checksum_value(const char* c, size_t size) {
        _cntl->_checksum_value.assign(c, size);
    }

    void set_checksum_value(const std::string& c) {
        _cntl->_checksum_value = c;
    }

    const std::string& checksum_value() const { return _cntl->_checksum_value; }

private:
    Controller* _cntl;
};

// Inherit this class to intercept Controller::IssueRPC. This is an internal
// utility only useable by brpc developers.
class RPCSender {
public:
    virtual ~RPCSender() {}
    virtual int IssueRPC(int64_t start_realtime_us) = 0;
};

} // namespace brpc


#endif // BRPC_CONTROLLER_PRIVATE_ACCESSOR_H