rocketmq-transport 1.0.0

// Copyright 2023 The RocketMQ Rust Authors
//
// 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.

//!
//! # Key Improvements
//!
//! - **Code Reuse**: Eliminated 210 lines of duplicated code using generic handler
//! - **Error Handling**: Preserved full error context with `thiserror` integration
//! - **Performance**: Removed unnecessary boxing operations
//! - **Async Pattern**: Added callback-based invocation for non-blocking scenarios
//!
//! # Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────────────┐
//! │                    RpcClientImpl                            │
//! ├─────────────────────────────────────────────────────────────┤
//! │                                                             │
//! │  ┌──────────────┐        ┌───────────────────────┐         │
//! │  │ Client Hooks │───────►│ Request Interceptor   │         │
//! │  └──────────────┘        └───────────────────────┘         │
//! │         │                          │                        │
//! │         ↓                          ↓                        │
//! │  ┌──────────────────────────────────────────────┐          │
//! │  │   Generic Handler (eliminates duplication)   │          │
//! │  │   handle_request<H, R>()                     │          │
//! │  └──────────────────────────────────────────────┘          │
//! │         │                                                   │
//! │         ↓                                                   │
//! │  ┌──────────────────────────────────────────────┐          │
//! │  │   RemotingClient (Connection Pool)           │          │
//! │  │   - invoke_request (sync)                    │          │
//! │  │   - invoke_request_oneway                    │          │
//! │  └──────────────────────────────────────────────┘          │
//! │                                                             │
//! └─────────────────────────────────────────────────────────────┘
//! ```

use std::any::Any;
use std::sync::Arc;

use cheetah_string::CheetahString;
use rocketmq_error::SharedError;
use rocketmq_model::common::message::message_queue::MessageQueue;
use rocketmq_runtime::TaskId;
use tracing::error;
use tracing::trace;

use crate::clients::rocketmq_tokio_client::TransportClient;
use crate::deadline::RequestDeadline;
use crate::error_helpers::rpc_broker_address_not_found;
use crate::error_helpers::rpc_request_failed;
use crate::error_helpers::rpc_request_unsupported;
use crate::error_helpers::rpc_response_failed;
use crate::request_outcome::OutboundRequestContractReason;
use crate::request_outcome::OutboundRequestOutcome;
use crate::request_outcome::OutboundRequestRejectionReason;
use crate::request_processor::default_request_processor::DefaultRequestProcessor;
use crate::rpc::client_metadata::ClientMetadata;
use crate::rpc::rpc_client::RpcClient;
use crate::rpc::rpc_client_hook::RpcClientHookFn;
use crate::rpc::rpc_client_utils::RpcClientUtils;
use crate::rpc::rpc_request::RpcRequest;
use crate::rpc::rpc_response::RpcResponse;
use rocketmq_protocol::code::request_code::RequestCode;
use rocketmq_protocol::code::response_code::ResponseCode;
use rocketmq_protocol::protocol::command_custom_header::CommandCustomHeader;
use rocketmq_protocol::protocol::command_custom_header::FromMap;
use rocketmq_protocol::protocol::header::get_earliest_msg_storetime_response_header::GetEarliestMsgStoretimeResponseHeader;
use rocketmq_protocol::protocol::header::get_max_offset_response_header::GetMaxOffsetResponseHeader;
use rocketmq_protocol::protocol::header::get_min_offset_response_header::GetMinOffsetResponseHeader;
use rocketmq_protocol::protocol::header::message_operation_header::TopicRequestHeaderTrait;
use rocketmq_protocol::protocol::header::pull_message_response_header::PullMessageResponseHeader;
use rocketmq_protocol::protocol::header::query_consumer_offset_response_header::QueryConsumerOffsetResponseHeader;
use rocketmq_protocol::protocol::header::search_offset_response_header::SearchOffsetResponseHeader;
use rocketmq_protocol::protocol::header::update_consumer_offset_header::UpdateConsumerOffsetResponseHeader;
use rocketmq_protocol::protocol::remoting_command::RemotingCommand;

/// Configuration for response handling
struct ResponseConfig {
    /// Expected successful response codes
    success_codes: &'static [ResponseCode],
}

impl ResponseConfig {
    const fn new(success_codes: &'static [ResponseCode]) -> Self {
        Self { success_codes }
    }
}

fn rpc_request_response(
    result: Result<OutboundRequestOutcome, crate::error::TransportError>,
) -> Result<RemotingCommand, rocketmq_error::SharedError> {
    match result {
        Ok(OutboundRequestOutcome::Response(response)) => Ok(response),
        Ok(OutboundRequestOutcome::Rejected(rejection)) => {
            let error = match rejection.reason() {
                OutboundRequestRejectionReason::DeadlineExpired => {
                    rocketmq_error::Error::new(&rocketmq_error::CORE_OPERATION_TIMED_OUT).with_context(
                        rocketmq_error::ErrorContext::new()
                            .with_text(rocketmq_error::fields::OPERATION_DIAGNOSTIC, "rpc_request")
                            .with_u64(
                                rocketmq_error::fields::TIMEOUT_MS,
                                rejection.timeout_millis().unwrap_or_default(),
                            ),
                    )
                }
                OutboundRequestRejectionReason::ClientStopping => {
                    rocketmq_error::Error::new(&rocketmq_error::CLIENT_LIFECYCLE_NOT_STARTED)
                }
                OutboundRequestRejectionReason::QueueSaturated => {
                    rocketmq_error::Error::new(&rocketmq_error::TRANSPORT_ADMISSION_QUEUE_SATURATED)
                }
                OutboundRequestRejectionReason::Cancelled
                | OutboundRequestRejectionReason::SessionClosed
                | OutboundRequestRejectionReason::EndpointUnavailable => {
                    rocketmq_error::Error::new(&rocketmq_error::TRANSPORT_CONNECTION_FAILED)
                }
            };
            Err(Arc::new(error))
        }
        Ok(OutboundRequestOutcome::Contract(contract)) => {
            let descriptor = match contract.reason() {
                OutboundRequestContractReason::NameServerEndpointMissing => {
                    &rocketmq_error::TRANSPORT_CONNECTION_FAILED
                }
            };
            Err(Arc::new(rocketmq_error::Error::new(descriptor)))
        }
        Err(error) => Err(error.into_shared_error()),
    }
}

/// High-performance RPC client implementation
///
/// # Thread Safety
///
/// This struct is `Send + Sync` safe and can be shared across threads using `Arc<RpcClientImpl>`.
///
/// # Example
///
/// ```rust,ignore
/// use crate::rpc::rpc_client_impl::RpcClientImpl;
///
/// let client = RpcClientImpl::new(client_metadata, remoting_client);
///
/// // Synchronous invocation
/// let response = client.invoke(request, 3000).await?;
///
/// // Callback-based invocation (non-blocking)
/// client.invoke_with_callback(request, 3000, |result| {
///     match result {
///         Ok(response) => println!("Success: {:?}", response),
///         Err(err) => eprintln!("Failed: {}", err),
///     }
/// });
/// ```
#[derive(Clone)]
pub struct RpcClientImpl {
    /// Client metadata for broker address resolution
    client_metadata: Arc<ClientMetadata>,
    /// Underlying remoting client for network communication
    remoting_client: Arc<TransportClient<DefaultRequestProcessor>>,
    /// Registered client hooks for request interception
    client_hook_list: Vec<RpcClientHookFn>,
}

impl RpcClientImpl {
    /// Creates a new RPC client instance
    ///
    /// # Arguments
    ///
    /// * `client_metadata` - Metadata for broker address resolution
    /// * `remoting_client` - Network client for sending requests
    pub fn new(
        client_metadata: Arc<ClientMetadata>,
        remoting_client: Arc<TransportClient<DefaultRequestProcessor>>,
    ) -> Self {
        RpcClientImpl {
            client_metadata,
            remoting_client,
            client_hook_list: Vec::new(),
        }
    }

    /// Registers a client hook for request/response interception
    ///
    /// Hooks are executed in registration order before each request.
    pub fn register_client_hook(&mut self, client_hook: RpcClientHookFn) {
        self.client_hook_list.push(client_hook);
    }

    /// Clears all registered client hooks
    pub fn clear_client_hook(&mut self) {
        self.client_hook_list.clear();
    }

    fn create_request_command<H>(
        &self,
        addr: &CheetahString,
        request: RpcRequest<H>,
        timeout_millis: u64,
    ) -> Result<(i32, RemotingCommand), SharedError>
    where
        H: CommandCustomHeader + TopicRequestHeaderTrait,
    {
        let request_code = request.code;
        let request_command = RpcClientUtils::try_create_command_for_rpc_request(request)
            .map_err(|err| rpc_request_failed(addr.to_string(), request_code, timeout_millis, err))?;
        Ok((request_code, request_command))
    }

    /// Resolves broker address by name, returning error if not found
    fn get_broker_addr_by_name(&self, broker_name: &str) -> Result<CheetahString, SharedError> {
        self.client_metadata
            .find_master_broker_addr(broker_name)
            .ok_or_else(|| rpc_broker_address_not_found(broker_name))
    }

    /// Generic request handler eliminating code duplication
    ///
    ///
    /// # Type Parameters
    ///
    /// * `H` - Request header type
    /// * `R` - Response header type
    ///
    /// # Arguments
    ///
    /// * `addr` - Target broker address
    /// * `request` - RPC request with custom header
    /// * `timeout_millis` - Request timeout in milliseconds
    /// * `config` - Response handling configuration
    async fn handle_request<H, R>(
        &self,
        addr: &CheetahString,
        request: RpcRequest<H>,
        deadline: RequestDeadline,
        config: ResponseConfig,
    ) -> Result<RpcResponse, SharedError>
    where
        H: CommandCustomHeader + TopicRequestHeaderTrait,
        R: CommandCustomHeader + FromMap<Target = R, Error = rocketmq_error::Error> + Send + Sync + 'static,
    {
        let timeout_millis = deadline.budget_millis();
        trace!(
            "Sending RPC request: addr={}, code={}, timeout={}ms",
            addr,
            request.code,
            timeout_millis
        );

        let (request_code, request_command) = self.create_request_command(addr, request, timeout_millis)?;
        deadline
            .ensure_before_send()
            .map_err(|err| rpc_request_failed(addr.to_string(), request_code, timeout_millis, err))?;

        let response = rpc_request_response(
            self.remoting_client
                .invoke_request_with_deadline(Some(addr), request_command, deadline)
                .await,
        )
        .map_err(|err| {
            error!(
                "RPC request failed: addr={}, code={}, error={}",
                addr, request_code, err
            );
            rpc_request_failed(addr.to_string(), request_code, timeout_millis, err)
        })?;

        let response_code = ResponseCode::from(response.code());

        if !config.success_codes.contains(&response_code) {
            return Err(rpc_response_failed(response.code()));
        }

        let response_header = response
            .decode_command_custom_header::<R>()
            .map_err(|err| rpc_request_failed(addr.to_string(), request_code, timeout_millis, err))?;

        let body = response.body().map(|value| Box::new(value.clone()) as Box<dyn Any>);

        Ok(RpcResponse::new(response.code(), Box::new(response_header), body))
    }

    /// Handles PULL_MESSAGE request with multiple success codes
    async fn handle_pull_message<H: CommandCustomHeader + TopicRequestHeaderTrait>(
        &self,
        addr: &CheetahString,
        request: RpcRequest<H>,
        deadline: RequestDeadline,
    ) -> Result<RpcResponse, SharedError> {
        const PULL_SUCCESS_CODES: &[ResponseCode] = &[
            ResponseCode::Success,
            ResponseCode::PullNotFound,
            ResponseCode::PullRetryImmediately,
            ResponseCode::PullOffsetMoved,
        ];

        self.handle_request::<H, PullMessageResponseHeader>(
            addr,
            request,
            deadline,
            ResponseConfig::new(PULL_SUCCESS_CODES),
        )
        .await
    }

    /// Handles QUERY_CONSUMER_OFFSET with QUERY_NOT_FOUND support
    async fn handle_query_consumer_offset<H: CommandCustomHeader + TopicRequestHeaderTrait>(
        &self,
        addr: &CheetahString,
        request: RpcRequest<H>,
        deadline: RequestDeadline,
    ) -> Result<RpcResponse, SharedError> {
        let timeout_millis = deadline.budget_millis();
        let (request_code, request_command) = self.create_request_command(addr, request, timeout_millis)?;
        deadline
            .ensure_before_send()
            .map_err(|err| rpc_request_failed(addr.to_string(), request_code, timeout_millis, err))?;

        let response = rpc_request_response(
            self.remoting_client
                .invoke_request_with_deadline(Some(addr), request_command, deadline)
                .await,
        )
        .map_err(|err| rpc_request_failed(addr.to_string(), request_code, timeout_millis, err))?;

        match ResponseCode::from(response.code()) {
            ResponseCode::Success => {
                let response_header = response
                    .decode_command_custom_header::<QueryConsumerOffsetResponseHeader>()
                    .map_err(|err| rpc_request_failed(addr.to_string(), request_code, timeout_millis, err))?;
                let body = response.body().map(|value| Box::new(value.clone()) as Box<dyn Any>);
                Ok(RpcResponse::new(response.code(), Box::new(response_header), body))
            }
            ResponseCode::QueryNotFound => {
                // Special case: no offset found (not an error)
                Ok(RpcResponse::new_option(response.code(), None))
            }
            _ => Err(rpc_response_failed(response.code())),
        }
    }

    /// Executes client hooks before request, short-circuiting if hook returns response
    fn execute_hooks<H: CommandCustomHeader + TopicRequestHeaderTrait>(
        &self,
        request: &RpcRequest<H>,
    ) -> Result<Option<RpcResponse>, rocketmq_error::SharedError> {
        for hook in &self.client_hook_list {
            if let Some(response) = hook(Some(&request.header), None)? {
                trace!("Request intercepted by client hook");
                return Ok(Some(response));
            }
        }
        Ok(None)
    }

    async fn invoke_until<H: CommandCustomHeader + TopicRequestHeaderTrait>(
        &self,
        request: RpcRequest<H>,
        deadline: RequestDeadline,
    ) -> Result<RpcResponse, rocketmq_error::SharedError> {
        let timeout_millis = deadline.budget_millis();
        if let Some(response) = self.execute_hooks(&request)? {
            if deadline.is_expired() {
                return Err(crate::error_helpers::response_timeout(timeout_millis));
            }
            return Ok(response);
        }

        let broker_name = request
            .header
            .broker_name()
            .ok_or_else(|| rpc_broker_address_not_found("<missing brokerName>"))?;
        let addr = self.get_broker_addr_by_name(broker_name.as_ref())?;
        deadline.ensure_before_send()?;

        let result = match RequestCode::from(request.code) {
            RequestCode::PullMessage => self.handle_pull_message(&addr, request, deadline).await?,
            RequestCode::GetMinOffset => {
                self.handle_request::<H, GetMinOffsetResponseHeader>(
                    &addr,
                    request,
                    deadline,
                    ResponseConfig::new(&[ResponseCode::Success]),
                )
                .await?
            }
            RequestCode::GetMaxOffset => {
                self.handle_request::<H, GetMaxOffsetResponseHeader>(
                    &addr,
                    request,
                    deadline,
                    ResponseConfig::new(&[ResponseCode::Success]),
                )
                .await?
            }
            RequestCode::SearchOffsetByTimestamp => {
                self.handle_request::<H, SearchOffsetResponseHeader>(
                    &addr,
                    request,
                    deadline,
                    ResponseConfig::new(&[ResponseCode::Success]),
                )
                .await?
            }
            RequestCode::GetEarliestMsgStoreTime => {
                self.handle_request::<H, GetEarliestMsgStoretimeResponseHeader>(
                    &addr,
                    request,
                    deadline,
                    ResponseConfig::new(&[ResponseCode::Success]),
                )
                .await?
            }
            RequestCode::QueryConsumerOffset => self.handle_query_consumer_offset(&addr, request, deadline).await?,
            RequestCode::UpdateConsumerOffset => {
                self.handle_request::<H, UpdateConsumerOffsetResponseHeader>(
                    &addr,
                    request,
                    deadline,
                    ResponseConfig::new(&[ResponseCode::Success]),
                )
                .await?
            }
            RequestCode::GetTopicStatsInfo | RequestCode::GetTopicConfig => {
                let (request_code, request_command) = self.create_request_command(&addr, request, timeout_millis)?;
                deadline.ensure_before_send()?;
                let response = rpc_request_response(
                    self.remoting_client
                        .invoke_request_with_deadline(Some(&addr), request_command, deadline)
                        .await,
                )
                .map_err(|err| rpc_request_failed(addr.to_string(), request_code, timeout_millis, err))?;

                if response.code() != ResponseCode::Success as i32 {
                    return Err(rpc_response_failed(response.code()));
                }

                let body = response.body().map(|value| Box::new(value.clone()) as Box<dyn Any>);
                RpcResponse::new_option(response.code(), body)
            }
            _ => return Err(rpc_request_unsupported(request.code)),
        };

        Ok(result)
    }
}

impl RpcClient for RpcClientImpl {
    /// Invokes RPC request synchronously
    ///
    /// # Flow
    ///
    /// 1. Execute client hooks (may short-circuit)
    /// 2. Resolve broker address
    /// 3. Dispatch to specific handler based on request code
    /// 4. Return response or error
    async fn invoke<H: CommandCustomHeader + TopicRequestHeaderTrait>(
        &self,
        request: RpcRequest<H>,
        timeout_millis: u64,
    ) -> Result<RpcResponse, rocketmq_error::SharedError> {
        self.invoke_until(request, RequestDeadline::from_timeout_millis(timeout_millis))
            .await
    }

    /// Invokes RPC request with message queue context
    ///
    /// This method resolves broker name from message queue and delegates to `invoke()`.
    async fn invoke_mq<H: CommandCustomHeader + TopicRequestHeaderTrait>(
        &self,
        mq: MessageQueue,
        mut request: RpcRequest<H>,
        timeout_millis: u64,
    ) -> Result<RpcResponse, rocketmq_error::SharedError> {
        let deadline = RequestDeadline::from_timeout_millis(timeout_millis);
        if let Some(broker_name) = self.client_metadata.get_broker_name_from_message_queue(&mq) {
            request.header.set_broker_name(broker_name);
        }
        self.invoke_until(request, deadline).await
    }
}

// ==================== Callback-based Async Extension ====================

impl RpcClientImpl {
    /// Invokes RPC request with callback (non-blocking)
    ///
    /// Unlike `invoke()` which awaits the response, this method spawns an async task
    /// and invokes the callback when the response arrives.
    ///
    /// # Use Cases
    ///
    /// - Batch requests where you don't want to block on each response
    /// - Pipeline processing where requests can be issued in parallel
    /// - Event-driven architectures
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// client.invoke_with_callback(request, 3000, |result| {
    ///     match result {
    ///         Ok(response) => process_response(response),
    ///         Err(err) => log_error(err),
    ///     }
    /// });
    /// ```
    pub fn invoke_with_callback<H, F>(&self, request: RpcRequest<H>, timeout_millis: u64, callback: F) -> TaskId
    where
        H: CommandCustomHeader + TopicRequestHeaderTrait + Send + 'static,
        F: FnOnce(Result<RpcResponse, rocketmq_error::SharedError>) + Send + 'static,
    {
        let deadline = RequestDeadline::from_timeout_millis(timeout_millis);
        let client_metadata = self.client_metadata.clone();
        let remoting_client = self.remoting_client.clone();
        let remoting_client_for_task = remoting_client.clone();
        let hooks = self.client_hook_list.clone();

        remoting_client
            .spawn_worker_task("remoting.rpc.callback", async move {
                // Create a temporary client instance for the async task
                let temp_client = RpcClientImpl {
                    client_metadata,
                    remoting_client: remoting_client_for_task,
                    client_hook_list: hooks,
                };

                let result = temp_client.invoke_until(request, deadline).await;
                callback(result);
            })
            .expect("RPC callback task should be spawned on the remoting worker task group")
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use super::*;
    use crate::runtime::config::client_config::TransportClientConfig;
    use rocketmq_protocol::protocol::header::get_min_offset_request_header::GetMinOffsetRequestHeader;

    fn test_service_context(name: &'static str) -> rocketmq_runtime::ChildServiceContext {
        rocketmq_runtime::RuntimeContext::from_current(name).service_context("rpc-client-service")
    }

    #[test]
    fn test_error_formatting() {
        let err = rpc_broker_address_not_found("broker-a");
        assert_eq!(err.descriptor(), &rocketmq_error::RPC_BROKER_ADDRESS_NOT_FOUND);

        let err = rpc_response_failed(1);
        assert_eq!(err.descriptor(), &rocketmq_error::RPC_RESPONSE_FAILED);
    }

    #[tokio::test]
    async fn invoke_without_broker_name_returns_typed_error_instead_of_panicking() {
        let client_metadata = Arc::new(ClientMetadata::new());
        let remoting_client = Arc::new(TransportClient::build_for_test(
            Arc::new(TransportClientConfig::default()),
            DefaultRequestProcessor,
            test_service_context("rpc-client-error-test"),
        ));
        let rpc_client = RpcClientImpl::new(client_metadata, remoting_client);
        let request = RpcRequest::new(
            RequestCode::GetMinOffset.into(),
            GetMinOffsetRequestHeader {
                topic: CheetahString::from_static_str("TopicA"),
                queue_id: 0,
                topic_request_header: None,
            },
            None,
        );

        let Err(error) = rpc_client.invoke(request, 3000).await else {
            panic!("missing brokerName should return a typed error");
        };

        assert_eq!(error.descriptor(), &rocketmq_error::RPC_BROKER_ADDRESS_NOT_FOUND);
    }

    #[tokio::test]
    async fn invoke_with_callback_uses_remoting_worker_task_group() {
        let client_metadata = Arc::new(ClientMetadata::new());
        let remoting_client = Arc::new(TransportClient::build_for_test(
            Arc::new(TransportClientConfig::default()),
            DefaultRequestProcessor,
            test_service_context("rpc-client-callback-test"),
        ));
        let rpc_client = RpcClientImpl::new(client_metadata, remoting_client.clone());
        let request = RpcRequest::new(
            RequestCode::GetMinOffset.into(),
            GetMinOffsetRequestHeader {
                topic: CheetahString::from_static_str("TopicA"),
                queue_id: 0,
                topic_request_header: None,
            },
            None,
        );
        let (tx, rx) = tokio::sync::oneshot::channel();

        let task_id = rpc_client.invoke_with_callback(request, 3000, move |result| {
            let _ = tx.send(result.is_err());
        });

        let worker_task_group = remoting_client.worker_task_group().expect("worker task group");
        assert!(
            worker_task_group
                .wait_task(task_id, std::time::Duration::from_secs(1))
                .await,
            "callback task should finish"
        );
        assert!(rx.await.expect("callback result should be sent"));
        assert_eq!(
            worker_task_group.lifecycle_state(),
            rocketmq_runtime::TaskGroupLifecycleState::Open
        );
        assert_eq!(worker_task_group.task_count(), 0);

        remoting_client.shutdown();

        assert_eq!(
            worker_task_group.lifecycle_state(),
            rocketmq_runtime::TaskGroupLifecycleState::ShutdownCompleted
        );
    }
}