Skip to main content

rocketmq_protocol/rpc/
rpc_response.rs

1// Copyright 2023 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8
9use std::any::Any;
10
11use rocketmq_error::Error;
12use rocketmq_error::ViewValueRef;
13use rocketmq_error::BROKER_OPERATION_FAILED;
14
15use crate::protocol::CommandCustomHeader;
16
17#[derive(Default)]
18pub struct RpcResponse {
19    pub code: i32,
20    pub header: Option<Box<dyn CommandCustomHeader + Send + Sync + 'static>>,
21    pub body: Option<Box<dyn Any>>,
22    pub exception: Option<Error>,
23}
24
25impl RpcResponse {
26    pub fn get_header<T>(&self) -> Option<&T>
27    where
28        T: CommandCustomHeader + Send + Sync + 'static,
29    {
30        self.header.as_ref()?.as_any().downcast_ref::<T>()
31    }
32
33    pub fn get_header_mut<T>(&mut self) -> Option<&mut T>
34    where
35        T: CommandCustomHeader + Send + Sync + 'static,
36    {
37        self.header.as_mut()?.as_any_mut().downcast_mut::<T>()
38    }
39
40    pub fn new_exception(exception: Option<Error>) -> Self {
41        Self {
42            code: exception.as_ref().map_or(0, broker_response_code),
43            header: None,
44            body: None,
45            exception,
46        }
47    }
48
49    pub fn new(
50        code: i32,
51        header: Box<dyn CommandCustomHeader + Send + Sync + 'static>,
52        body: Option<Box<dyn Any>>,
53    ) -> Self {
54        Self {
55            code,
56            header: Some(header),
57            body,
58            exception: None,
59        }
60    }
61
62    pub fn new_option(code: i32, body: Option<Box<dyn Any>>) -> Self {
63        Self {
64            code,
65            header: None,
66            body,
67            exception: None,
68        }
69    }
70}
71
72fn broker_response_code(error: &Error) -> i32 {
73    if error.descriptor() != &BROKER_OPERATION_FAILED {
74        return 0;
75    }
76
77    error
78        .diagnostic_view()
79        .ok()
80        .and_then(|view| {
81            view.fields()
82                .find(|field| field.name() == rocketmq_error::fields::BROKER_CODE.schema().name())
83                .and_then(|field| match field.value() {
84                    ViewValueRef::I64(code) => i32::try_from(code).ok(),
85                    _ => None,
86                })
87        })
88        .unwrap_or(0)
89}
90
91#[cfg(test)]
92mod tests {
93    use rocketmq_error::fields;
94    use rocketmq_error::ErrorContext;
95
96    use super::*;
97    use crate::protocol::header::empty_header::EmptyHeader;
98    use crate::protocol::header::get_all_topic_config_response_header::GetAllTopicConfigResponseHeader;
99
100    #[test]
101    fn exception_response_uses_zero_without_a_valid_broker_code() {
102        let empty = RpcResponse::new_exception(None);
103        assert_eq!(empty.code, 0);
104        assert!(empty.exception.is_none());
105        assert!(empty.header.is_none());
106        assert!(empty.body.is_none());
107
108        for error in [
109            crate::error::invalid_argument("invalid"),
110            Error::new(&BROKER_OPERATION_FAILED),
111        ] {
112            let descriptor = error.descriptor();
113            let response = RpcResponse::new_exception(Some(error));
114            assert_eq!(response.code, 0);
115            assert_eq!(response.exception.as_ref().unwrap().descriptor(), descriptor);
116        }
117    }
118
119    #[test]
120    fn exception_response_checks_broker_code_bounds_without_truncation() {
121        for (code, expected) in [
122            (i64::from(i32::MIN), i32::MIN),
123            (i64::from(i32::MAX), i32::MAX),
124            (i64::from(i32::MIN) - 1, 0),
125            (i64::from(i32::MAX) + 1, 0),
126        ] {
127            let error = Error::new(&BROKER_OPERATION_FAILED)
128                .with_context(ErrorContext::new().with_i64(fields::BROKER_CODE, code));
129            let response = RpcResponse::new_exception(Some(error));
130            assert_eq!(response.code, expected);
131            assert_eq!(
132                response.exception.as_ref().unwrap().descriptor(),
133                &BROKER_OPERATION_FAILED
134            );
135        }
136    }
137
138    #[test]
139    fn typed_headers_support_checked_reads_and_mutation_without_changing_body() {
140        let header = GetAllTopicConfigResponseHeader {
141            total_topic_num: Some(3),
142        };
143        let mut response = RpcResponse::new(17, Box::new(header), Some(Box::new(vec![1_u8, 2])));
144        assert_eq!(response.code, 17);
145        assert!(response.exception.is_none());
146        assert!(response.get_header::<EmptyHeader>().is_none());
147        assert!(response.get_header_mut::<EmptyHeader>().is_none());
148        assert_eq!(
149            response
150                .get_header::<GetAllTopicConfigResponseHeader>()
151                .unwrap()
152                .total_topic_num,
153            Some(3)
154        );
155        response
156            .get_header_mut::<GetAllTopicConfigResponseHeader>()
157            .unwrap()
158            .total_topic_num = Some(5);
159        assert_eq!(
160            response
161                .get_header::<GetAllTopicConfigResponseHeader>()
162                .unwrap()
163                .total_topic_num,
164            Some(5)
165        );
166        assert_eq!(
167            response.body.as_ref().unwrap().downcast_ref::<Vec<u8>>().unwrap(),
168            &vec![1, 2]
169        );
170
171        let mut without_header = RpcResponse::new_option(23, Some(Box::new("body")));
172        assert_eq!(without_header.code, 23);
173        assert!(without_header.exception.is_none());
174        assert!(without_header.get_header::<EmptyHeader>().is_none());
175        assert!(without_header.get_header_mut::<EmptyHeader>().is_none());
176        assert_eq!(
177            without_header.body.as_ref().unwrap().downcast_ref::<&str>(),
178            Some(&"body")
179        );
180        assert!(RpcResponse::new_option(0, None).body.is_none());
181        assert!(RpcResponse::new(0, Box::new(EmptyHeader::default()), None)
182            .body
183            .is_none());
184    }
185
186    #[test]
187    fn exception_response_retains_canonical_broker_code() {
188        let error = Error::new(&BROKER_OPERATION_FAILED).with_context(
189            ErrorContext::new()
190                .with_text(fields::OPERATION_DIAGNOSTIC, "test")
191                .with_i64(fields::BROKER_CODE, 207)
192                .with_secret_presence(fields::MESSAGE_PRESENT),
193        );
194
195        let response = RpcResponse::new_exception(Some(error));
196
197        assert_eq!(response.code, 207);
198        assert_eq!(
199            response.exception.as_ref().map(Error::descriptor),
200            Some(&BROKER_OPERATION_FAILED)
201        );
202    }
203}