rocketmq_protocol/protocol/remoting_command/
custom_header.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//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use super::RemotingCommand;
18use crate::protocol::command_custom_header::CommandCustomHeader;
19use crate::protocol::command_custom_header::FromMap;
20use crate::protocol::command_custom_header::HeaderEncodeCapability;
21use crate::protocol::header_field_merge::merge_header_and_dynamic;
22use crate::ProtocolContractViolation;
23
24impl RemotingCommand {
25    pub fn set_command_custom_header<T>(mut self, command_custom_header: T) -> Self
26    where
27        T: CommandCustomHeader + Sync + Send + 'static,
28    {
29        self.invalidate_materialized_custom_header();
30        self.command_custom_header = Some(Arc::new(Box::new(command_custom_header)));
31        self.custom_header_to_net = false;
32        self
33    }
34
35    pub fn set_command_custom_header_boxed(
36        mut self,
37        command_custom_header: Box<dyn CommandCustomHeader + Send + Sync + 'static>,
38    ) -> Self {
39        self.invalidate_materialized_custom_header();
40        self.command_custom_header = Some(Arc::new(command_custom_header));
41        self.custom_header_to_net = false;
42        self
43    }
44
45    pub fn set_command_custom_header_origin<T>(mut self, command_custom_header: Option<T>) -> Self
46    where
47        T: std::ops::Deref<Target = Box<dyn CommandCustomHeader + Send + Sync + 'static>>,
48    {
49        self.invalidate_materialized_custom_header();
50        if let Some(header_fields) = command_custom_header.as_ref().and_then(|header| header.to_map()) {
51            self.ext_fields.get_or_insert_map().extend(header_fields);
52        }
53        self.command_custom_header = None;
54        self.custom_header_to_net = true;
55        self
56    }
57
58    pub fn set_command_custom_header_ref<T>(&mut self, command_custom_header: T)
59    where
60        T: CommandCustomHeader + Sync + Send + 'static,
61    {
62        self.invalidate_materialized_custom_header();
63        self.command_custom_header = Some(Arc::new(Box::new(command_custom_header)));
64        self.custom_header_to_net = false;
65    }
66
67    /// Fallibly merges the custom header into dynamic extension fields.
68    ///
69    /// # Errors
70    ///
71    /// Returns a typed validation, conversion, alias, or dynamic-field
72    /// collision error without mutating this command.
73    pub fn try_make_custom_header_to_net(&mut self) -> Result<(), ProtocolContractViolation> {
74        if self.custom_header_to_net {
75            return Ok(());
76        }
77
78        if let Some(header) = self.command_custom_header_ref() {
79            let merged = merge_header_and_dynamic(header, self.ext_fields.as_map())?;
80            self.ext_fields.replace_map(merged);
81        }
82        self.custom_header_to_net = true;
83        Ok(())
84    }
85
86    fn invalidate_materialized_custom_header(&mut self) {
87        if !self.custom_header_to_net {
88            return;
89        }
90
91        let owned_keys = match (self.command_custom_header_ref(), self.ext_fields.as_map()) {
92            (Some(header), Some(fields)) => {
93                let mut keys = fields
94                    .keys()
95                    .filter(|key| header.contains_wire_key(key.as_str()))
96                    .cloned()
97                    .collect::<Vec<_>>();
98                if let Some(legacy_fields) = header.to_map() {
99                    keys.extend(legacy_fields.into_keys());
100                }
101                keys
102            }
103            _ => Vec::new(),
104        };
105        if let Some(fields) = self.ext_fields.as_map_mut() {
106            for key in owned_keys {
107                fields.remove(&key);
108            }
109        }
110        self.custom_header_to_net = false;
111    }
112
113    pub fn decode_command_custom_header<T>(&self) -> rocketmq_error::Result<T>
114    where
115        T: FromMap<Target = T, Error = rocketmq_error::Error>,
116    {
117        if T::SUPPORTS_HEADER_FIELD_SOURCE {
118            if let Some(source) = self.ext_fields.as_field_source() {
119                return T::from_field_source(source);
120            }
121        }
122        match self.ext_fields.as_map() {
123            None => Err(crate::error::serialization_decode_failed(
124                "header",
125                "extension fields absent",
126            )),
127            Some(header) => T::from(header),
128        }
129    }
130
131    pub fn decode_command_custom_header_fast<T>(&self) -> rocketmq_error::Result<T>
132    where
133        T: FromMap<Target = T, Error = rocketmq_error::Error>,
134        T: Default + CommandCustomHeader,
135    {
136        if T::SUPPORTS_HEADER_FIELD_SOURCE {
137            if let Some(source) = self.ext_fields.as_field_source() {
138                return T::from_field_source(source);
139            }
140        }
141        match self.ext_fields.as_map() {
142            None => Err(crate::error::serialization_decode_failed(
143                "header",
144                "extension fields absent",
145            )),
146            Some(header) => {
147                let mut target = T::default();
148                if target.support_fast_codec() {
149                    target.decode_fast(header)?;
150                    target.check_fields()?;
151                    Ok(target)
152                } else {
153                    T::from(header)
154                }
155            }
156        }
157    }
158
159    /// Decodes a required custom request header and classifies any failure at
160    /// the request-header boundary.
161    ///
162    /// `operation` must be a static, low-cardinality description. It is exposed
163    /// as structured error context, while the source retains the decoder cause.
164    ///
165    /// # Errors
166    ///
167    /// Returns a canonical request-header error when the extension fields are
168    /// absent or the header cannot be decoded.
169    pub fn decode_required_header<T>(&self, operation: &'static str) -> rocketmq_error::Result<T>
170    where
171        T: FromMap<Target = T, Error = rocketmq_error::Error>,
172    {
173        self.decode_command_custom_header::<T>()
174            .map_err(|source| required_header_decode_error(operation, source))
175    }
176
177    /// Decodes a required custom request header through the fast codec when the
178    /// header supports it and classifies any failure at the request-header
179    /// boundary.
180    ///
181    /// `operation` must be a static, low-cardinality description. It is exposed
182    /// as structured error context, while the source retains the decoder cause.
183    ///
184    /// # Errors
185    ///
186    /// Returns a canonical request-header error when the extension fields are
187    /// absent or the header cannot be decoded.
188    pub fn decode_required_header_fast<T>(&self, operation: &'static str) -> rocketmq_error::Result<T>
189    where
190        T: FromMap<Target = T, Error = rocketmq_error::Error>,
191        T: Default + CommandCustomHeader,
192    {
193        self.decode_command_custom_header_fast::<T>()
194            .map_err(|source| required_header_decode_error(operation, source))
195    }
196
197    pub fn try_read_custom_header_ref<T>(&self) -> rocketmq_error::Result<&T>
198    where
199        T: CommandCustomHeader + Sync + Send + 'static,
200    {
201        match self.command_custom_header.as_ref() {
202            None => Err(Self::custom_header_missing_error::<T>()),
203            Some(value) => value
204                .as_ref()
205                .as_any()
206                .downcast_ref::<T>()
207                .ok_or_else(Self::custom_header_type_mismatch_error::<T>),
208        }
209    }
210
211    pub fn try_read_custom_header_mut<T>(&mut self) -> rocketmq_error::Result<&mut T>
212    where
213        T: CommandCustomHeader + Sync + Send + 'static,
214    {
215        match self.command_custom_header.as_ref() {
216            None => return Err(Self::custom_header_missing_error::<T>()),
217            Some(value) if Arc::strong_count(value) != 1 => return Err(Self::custom_header_shared_error()),
218            Some(value) if !value.as_ref().as_any().is::<T>() => {
219                return Err(Self::custom_header_type_mismatch_error::<T>());
220            }
221            Some(_) => {}
222        }
223        self.invalidate_materialized_custom_header();
224        Arc::get_mut(
225            self.command_custom_header
226                .as_mut()
227                .ok_or_else(Self::custom_header_missing_error::<T>)?,
228        )
229        .ok_or_else(Self::custom_header_shared_error)?
230        .as_mut()
231        .as_any_mut()
232        .downcast_mut::<T>()
233        .ok_or_else(Self::custom_header_type_mismatch_error::<T>)
234    }
235
236    pub fn command_custom_header_ref(&self) -> Option<&dyn CommandCustomHeader> {
237        match self.command_custom_header.as_ref() {
238            None => None,
239            Some(value) => Some(value.as_ref().as_ref()),
240        }
241    }
242
243    pub(crate) fn custom_header_encode_capability(&self) -> HeaderEncodeCapability {
244        if self.custom_header_to_net {
245            HeaderEncodeCapability::MapOnly
246        } else {
247            self.command_custom_header_ref()
248                .map_or(HeaderEncodeCapability::MapOnly, CommandCustomHeader::encode_capability)
249        }
250    }
251
252    pub fn command_custom_header_mut(&mut self) -> Option<&mut dyn CommandCustomHeader> {
253        if self
254            .command_custom_header
255            .as_ref()
256            .is_none_or(|header| Arc::strong_count(header) != 1)
257        {
258            return None;
259        }
260        self.invalidate_materialized_custom_header();
261        match self.command_custom_header.as_mut() {
262            None => None,
263            Some(value) => Arc::get_mut(value).map(|header| header.as_mut() as &mut dyn CommandCustomHeader),
264        }
265    }
266
267    fn custom_header_missing_error<T>() -> rocketmq_error::Error
268    where
269        T: CommandCustomHeader + Sync + Send + 'static,
270    {
271        crate::error::serialization_decode_failed(
272            "header",
273            format!(
274                "Command custom header is missing; expected {}.",
275                std::any::type_name::<T>()
276            ),
277        )
278    }
279
280    fn custom_header_type_mismatch_error<T>() -> rocketmq_error::Error
281    where
282        T: CommandCustomHeader + Sync + Send + 'static,
283    {
284        crate::error::serialization_decode_failed(
285            "header",
286            format!(
287                "Command custom header type mismatch; expected {}.",
288                std::any::type_name::<T>()
289            ),
290        )
291    }
292
293    fn custom_header_shared_error() -> rocketmq_error::Error {
294        crate::error::serialization_decode_failed(
295            "header",
296            "Command custom header is shared by a cloned command and cannot be mutated safely.",
297        )
298    }
299}
300
301#[inline]
302fn required_header_decode_error(operation: &'static str, source: rocketmq_error::Error) -> rocketmq_error::Error {
303    crate::error::request_header_source(operation, source)
304}
rocketmq_protocol/protocol/remoting_command/custom_header.rs

rocketmq_protocol/protocol/remoting_command/
custom_header.rs
