common_meta/heartbeat/
handler.rs

1// Copyright 2023 Greptime Team
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 api::v1::meta::HeartbeatResponse;
18use async_trait::async_trait;
19use common_telemetry::error;
20
21use crate::error::Result;
22use crate::heartbeat::mailbox::{IncomingMessage, MailboxRef};
23
24pub mod invalidate_table_cache;
25pub mod parse_mailbox_message;
26pub mod suspend;
27#[cfg(test)]
28mod tests;
29
30pub type HeartbeatResponseHandlerExecutorRef = Arc<dyn HeartbeatResponseHandlerExecutor>;
31pub type HeartbeatResponseHandlerRef = Arc<dyn HeartbeatResponseHandler>;
32
33pub struct HeartbeatResponseHandlerContext {
34    pub mailbox: MailboxRef,
35    pub response: HeartbeatResponse,
36    pub incoming_message: Option<IncomingMessage>,
37}
38
39/// HandleControl
40///
41/// Controls process of handling heartbeat response.
42#[derive(Debug, PartialEq)]
43pub enum HandleControl {
44    Continue,
45    Done,
46}
47
48impl HeartbeatResponseHandlerContext {
49    pub fn new(mailbox: MailboxRef, response: HeartbeatResponse) -> Self {
50        Self {
51            mailbox,
52            response,
53            incoming_message: None,
54        }
55    }
56}
57
58/// HeartbeatResponseHandler
59///
60/// [`HeartbeatResponseHandler::is_acceptable`] returns true if handler can handle incoming [`HeartbeatResponseHandlerContext`].
61///
62/// [`HeartbeatResponseHandler::handle`] handles all or part of incoming [`HeartbeatResponseHandlerContext`].
63#[async_trait]
64pub trait HeartbeatResponseHandler: Send + Sync {
65    fn is_acceptable(&self, ctx: &HeartbeatResponseHandlerContext) -> bool;
66
67    async fn handle(&self, ctx: &mut HeartbeatResponseHandlerContext) -> Result<HandleControl>;
68}
69
70#[async_trait]
71pub trait HeartbeatResponseHandlerExecutor: Send + Sync {
72    async fn handle(&self, ctx: HeartbeatResponseHandlerContext) -> Result<()>;
73}
74
75pub struct HandlerGroupExecutor {
76    handlers: Vec<HeartbeatResponseHandlerRef>,
77}
78
79impl HandlerGroupExecutor {
80    pub fn new(handlers: Vec<HeartbeatResponseHandlerRef>) -> Self {
81        Self { handlers }
82    }
83}
84
85#[async_trait]
86impl HeartbeatResponseHandlerExecutor for HandlerGroupExecutor {
87    async fn handle(&self, mut ctx: HeartbeatResponseHandlerContext) -> Result<()> {
88        for handler in &self.handlers {
89            if !handler.is_acceptable(&ctx) {
90                continue;
91            }
92
93            match handler.handle(&mut ctx).await {
94                Ok(HandleControl::Done) => break,
95                Ok(HandleControl::Continue) => {}
96                Err(e) => {
97                    let mailbox_message_id = ctx
98                        .response
99                        .mailbox_message
100                        .as_ref()
101                        .map(|message| message.id);
102                    let json_payload_len =
103                        ctx.response.mailbox_message.as_ref().and_then(|message| {
104                            message.payload.as_ref().map(|payload| match payload {
105                                api::v1::meta::mailbox_message::Payload::Json(json) => json.len(),
106                            })
107                        });
108                    error!(
109                        %e;
110                        "Error while handling heartbeat response: mailbox_message_id={mailbox_message_id:?}, json_payload_len={json_payload_len:?}"
111                    );
112                    break;
113                }
114            }
115        }
116        Ok(())
117    }
118}
119
120#[cfg(test)]
121mod error_log_tests {
122    use std::fmt::Debug;
123    use std::sync::{Arc, Mutex};
124
125    use api::v1::meta::mailbox_message::Payload;
126    use api::v1::meta::{HeartbeatResponse, MailboxMessage};
127    use common_telemetry::tracing::field::{Field, Visit};
128    use common_telemetry::tracing::{Event, Subscriber};
129    use common_telemetry::tracing_subscriber::layer::{Context, SubscriberExt};
130    use common_telemetry::{tracing, tracing_subscriber};
131
132    use super::parse_mailbox_message::ParseMailboxMessageHandler;
133    use super::*;
134    use crate::heartbeat::mailbox::HeartbeatMailbox;
135
136    #[derive(Clone, Default)]
137    struct LogCapture(Arc<Mutex<Vec<String>>>);
138
139    impl<S> tracing_subscriber::Layer<S> for LogCapture
140    where
141        S: Subscriber,
142    {
143        fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
144            let mut visitor = FieldVisitor::default();
145            event.record(&mut visitor);
146            self.0.lock().unwrap().push(visitor.fields.join(", "));
147        }
148    }
149
150    #[derive(Default)]
151    struct FieldVisitor {
152        fields: Vec<String>,
153    }
154
155    impl Visit for FieldVisitor {
156        fn record_debug(&mut self, field: &Field, value: &dyn Debug) {
157            self.fields.push(format!("{}={value:?}", field.name()));
158        }
159
160        fn record_str(&mut self, field: &Field, value: &str) {
161            self.fields.push(format!("{}={value:?}", field.name()));
162        }
163    }
164
165    #[tokio::test(flavor = "current_thread")]
166    async fn test_handler_error_log_excludes_mailbox_payload() {
167        let payload_sentinel = "MALFORMED_PACKED_PAYLOAD_MUST_NOT_BE_LOGGED".repeat(256);
168        let payload = format!(
169            r#"{{"PackedGcRegions":{{"regions":[],"packed_file_refs_manifest":"{}"}}"#,
170            payload_sentinel
171        );
172        let payload_len = payload.len();
173        let capture = LogCapture::default();
174        let subscriber = tracing_subscriber::registry().with(capture.clone());
175        let _guard = tracing::subscriber::set_default(subscriber);
176        let (mailbox_tx, _) = tokio::sync::mpsc::channel(1);
177        let ctx = HeartbeatResponseHandlerContext::new(
178            Arc::new(HeartbeatMailbox::new(mailbox_tx)),
179            HeartbeatResponse {
180                mailbox_message: Some(MailboxMessage {
181                    id: 42,
182                    subject: "unsafe subject".to_string(),
183                    to: "unsafe recipient".to_string(),
184                    from: "unsafe sender".to_string(),
185                    payload: Some(Payload::Json(payload)),
186                    ..Default::default()
187                }),
188                ..Default::default()
189            },
190        );
191        let executor = HandlerGroupExecutor::new(vec![Arc::new(ParseMailboxMessageHandler)]);
192
193        executor.handle(ctx).await.unwrap();
194
195        let logs = capture.0.lock().unwrap().join("\n");
196        assert!(logs.contains("Error while handling heartbeat response"));
197        assert!(logs.contains("mailbox_message_id=Some(42)"));
198        assert!(logs.contains(&format!("json_payload_len=Some({payload_len})")));
199        assert!(!logs.contains(&payload_sentinel));
200    }
201}
common_meta/heartbeat/handler.rs

common_meta/heartbeat/
handler.rs
