1use std::collections::HashSet;
16use std::str::FromStr;
17
18use api::v1::meta::{DatanodeWorkloads, HeartbeatRequest, RequestHeader};
19use common_time::util as time_util;
20use lazy_static::lazy_static;
21use regex::Regex;
22use serde::{Deserialize, Serialize};
23use snafu::{OptionExt, ResultExt, ensure};
24use store_api::region_engine::{RegionRole, RegionStatistic};
25use store_api::storage::RegionId;
26use table::metadata::TableId;
27
28use crate::error::{self, DeserializeFromJsonSnafu, Result};
29use crate::heartbeat::utils::get_datanode_workloads;
30
31const DATANODE_STAT_PREFIX: &str = "__meta_datanode_stat";
32
33pub const REGION_STATISTIC_KEY: &str = "__region_statistic";
34
35lazy_static! {
36 pub(crate) static ref DATANODE_LEASE_KEY_PATTERN: Regex =
37 Regex::new("^__meta_datanode_lease-([0-9]+)-([0-9]+)$").unwrap();
38 static ref DATANODE_STAT_KEY_PATTERN: Regex =
39 Regex::new(&format!("^{DATANODE_STAT_PREFIX}-([0-9]+)-([0-9]+)$")).unwrap();
40 static ref INACTIVE_REGION_KEY_PATTERN: Regex =
41 Regex::new("^__meta_inactive_region-([0-9]+)-([0-9]+)-([0-9]+)$").unwrap();
42}
43
44#[derive(Debug, Clone, Default, Serialize, Deserialize)]
48pub struct Stat {
49 pub timestamp_millis: i64,
50 pub id: u64,
52 pub addr: String,
54 pub rcus: i64,
56 pub wcus: i64,
58 pub region_num: u64,
60 pub region_stats: Vec<RegionStat>,
62 pub topic_stats: Vec<TopicStat>,
64 pub node_epoch: u64,
66 pub datanode_workloads: DatanodeWorkloads,
68 pub gc_stat: Option<GcStat>,
70}
71
72#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
74pub struct RegionStat {
75 pub id: RegionId,
77 pub rcus: i64,
79 pub wcus: i64,
81 pub approximate_bytes: u64,
83 pub engine: String,
85 pub role: RegionRole,
87 pub num_rows: u64,
89 pub memtable_size: u64,
91 pub manifest_size: u64,
93 pub sst_size: u64,
95 pub sst_num: u64,
97 pub index_size: u64,
99 pub region_manifest: RegionManifestInfo,
101 pub written_bytes: u64,
103 pub data_topic_latest_entry_id: u64,
106 pub metadata_topic_latest_entry_id: u64,
110}
111
112#[derive(Debug, Clone, Serialize, Deserialize)]
113pub struct TopicStat {
114 pub topic: String,
116 pub latest_entry_id: u64,
118 pub record_size: u64,
120 pub record_num: u64,
122}
123
124pub trait TopicStatsReporter: Send + Sync {
126 fn reportable_topics(&mut self) -> Vec<TopicStat>;
128}
129
130#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
131pub enum RegionManifestInfo {
132 Mito {
133 manifest_version: u64,
134 flushed_entry_id: u64,
135 file_removed_cnt: u64,
137 },
138 Metric {
139 data_manifest_version: u64,
140 data_flushed_entry_id: u64,
141 metadata_manifest_version: u64,
142 metadata_flushed_entry_id: u64,
143 },
144}
145
146impl Stat {
147 #[inline]
148 pub fn is_empty(&self) -> bool {
149 self.region_stats.is_empty()
150 }
151
152 pub fn stat_key(&self) -> DatanodeStatKey {
153 DatanodeStatKey { node_id: self.id }
154 }
155
156 pub fn regions(&self) -> Vec<(RegionId, RegionRole)> {
158 self.region_stats.iter().map(|s| (s.id, s.role)).collect()
159 }
160
161 pub fn table_ids(&self) -> HashSet<TableId> {
163 self.region_stats.iter().map(|s| s.id.table_id()).collect()
164 }
165
166 pub fn retain_active_region_stats(&mut self, inactive_region_ids: &HashSet<RegionId>) {
168 if inactive_region_ids.is_empty() {
169 return;
170 }
171
172 self.region_stats
173 .retain(|r| !inactive_region_ids.contains(&r.id));
174 self.rcus = self.region_stats.iter().map(|s| s.rcus).sum();
175 self.wcus = self.region_stats.iter().map(|s| s.wcus).sum();
176 self.region_num = self.region_stats.len() as u64;
177 }
178
179 pub fn memory_size(&self) -> usize {
180 std::mem::size_of::<i64>() * 3 +
182 std::mem::size_of::<u64>() * 3 +
184 std::mem::size_of::<String>() + self.addr.capacity() +
186 self.region_stats.iter().map(|s| s.memory_size()).sum::<usize>()
188 }
189}
190
191impl RegionStat {
192 pub fn memory_size(&self) -> usize {
193 std::mem::size_of::<RegionRole>() +
195 std::mem::size_of::<RegionId>() +
197 std::mem::size_of::<i64>() * 4 +
199 std::mem::size_of::<u64>() * 5 +
201 std::mem::size_of::<String>() + self.engine.capacity() +
203 self.region_manifest.memory_size()
205 }
206}
207
208impl RegionManifestInfo {
209 pub fn memory_size(&self) -> usize {
210 match self {
211 RegionManifestInfo::Mito { .. } => std::mem::size_of::<u64>() * 2,
212 RegionManifestInfo::Metric { .. } => std::mem::size_of::<u64>() * 4,
213 }
214 }
215}
216
217impl TryFrom<&HeartbeatRequest> for Stat {
218 type Error = Option<RequestHeader>;
219
220 fn try_from(value: &HeartbeatRequest) -> std::result::Result<Self, Self::Error> {
221 let HeartbeatRequest {
222 header,
223 peer,
224 region_stats,
225 node_epoch,
226 node_workloads,
227 topic_stats,
228 extensions,
229 ..
230 } = value;
231
232 match (header, peer) {
233 (Some(header), Some(peer)) => {
234 let region_stats = region_stats
235 .iter()
236 .map(RegionStat::from)
237 .collect::<Vec<_>>();
238 let topic_stats = topic_stats.iter().map(TopicStat::from).collect::<Vec<_>>();
239
240 let datanode_workloads = get_datanode_workloads(node_workloads.as_ref());
241
242 let gc_stat = GcStat::from_extensions(extensions).map_err(|err| {
243 common_telemetry::error!(
244 "Failed to deserialize GcStat from extensions: {}",
245 err
246 );
247 header.clone()
248 })?;
249 Ok(Self {
250 timestamp_millis: time_util::current_time_millis(),
251 id: peer.id,
253 addr: peer.addr.clone(),
255 rcus: region_stats.iter().map(|s| s.rcus).sum(),
256 wcus: region_stats.iter().map(|s| s.wcus).sum(),
257 region_num: region_stats.len() as u64,
258 region_stats,
259 topic_stats,
260 node_epoch: *node_epoch,
261 datanode_workloads,
262 gc_stat,
263 })
264 }
265 (header, _) => Err(header.clone()),
266 }
267 }
268}
269
270impl From<store_api::region_engine::RegionManifestInfo> for RegionManifestInfo {
271 fn from(value: store_api::region_engine::RegionManifestInfo) -> Self {
272 match value {
273 store_api::region_engine::RegionManifestInfo::Mito {
274 manifest_version,
275 flushed_entry_id,
276 file_removed_cnt,
277 } => RegionManifestInfo::Mito {
278 manifest_version,
279 flushed_entry_id,
280 file_removed_cnt,
281 },
282 store_api::region_engine::RegionManifestInfo::Metric {
283 data_manifest_version,
284 data_flushed_entry_id,
285 metadata_manifest_version,
286 metadata_flushed_entry_id,
287 } => RegionManifestInfo::Metric {
288 data_manifest_version,
289 data_flushed_entry_id,
290 metadata_manifest_version,
291 metadata_flushed_entry_id,
292 },
293 }
294 }
295}
296
297impl From<&api::v1::meta::RegionStat> for RegionStat {
298 fn from(value: &api::v1::meta::RegionStat) -> Self {
299 let region_stat = value
300 .extensions
301 .get(REGION_STATISTIC_KEY)
302 .and_then(|value| RegionStatistic::deserialize_from_slice(value))
303 .unwrap_or_default();
304
305 Self {
306 id: RegionId::from_u64(value.region_id),
307 rcus: value.rcus,
308 wcus: value.wcus,
309 approximate_bytes: value.approximate_bytes as u64,
310 engine: value.engine.clone(),
311 role: RegionRole::from(value.role()),
312 num_rows: region_stat.num_rows,
313 memtable_size: region_stat.memtable_size,
314 manifest_size: region_stat.manifest_size,
315 sst_size: region_stat.sst_size,
316 sst_num: region_stat.sst_num,
317 index_size: region_stat.index_size,
318 region_manifest: region_stat.manifest.into(),
319 written_bytes: region_stat.written_bytes,
320 data_topic_latest_entry_id: region_stat.data_topic_latest_entry_id,
321 metadata_topic_latest_entry_id: region_stat.metadata_topic_latest_entry_id,
322 }
323 }
324}
325
326impl From<&api::v1::meta::TopicStat> for TopicStat {
327 fn from(value: &api::v1::meta::TopicStat) -> Self {
328 Self {
329 topic: value.topic_name.clone(),
330 latest_entry_id: value.latest_entry_id,
331 record_size: value.record_size,
332 record_num: value.record_num,
333 }
334 }
335}
336
337#[derive(Debug, Clone, Serialize, Deserialize, Default)]
338pub struct GcStat {
339 pub running_gc_tasks: u32,
341 pub gc_concurrency: u32,
343}
344
345impl GcStat {
346 pub const GC_STAT_KEY: &str = "__gc_stat";
347
348 pub fn new(running_gc_tasks: u32, gc_concurrency: u32) -> Self {
349 Self {
350 running_gc_tasks,
351 gc_concurrency,
352 }
353 }
354
355 pub fn into_extensions(&self, extensions: &mut std::collections::HashMap<String, Vec<u8>>) {
356 let bytes = serde_json::to_vec(self).unwrap_or_default();
357 extensions.insert(Self::GC_STAT_KEY.to_string(), bytes);
358 }
359
360 pub fn from_extensions(
361 extensions: &std::collections::HashMap<String, Vec<u8>>,
362 ) -> Result<Option<Self>> {
363 extensions
364 .get(Self::GC_STAT_KEY)
365 .map(|bytes| {
366 serde_json::from_slice(bytes).with_context(|_| DeserializeFromJsonSnafu {
367 input: String::from_utf8_lossy(bytes).to_string(),
368 })
369 })
370 .transpose()
371 }
372}
373
374#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
378pub struct DatanodeStatKey {
379 pub node_id: u64,
380}
381
382impl DatanodeStatKey {
383 pub fn prefix_key() -> Vec<u8> {
385 format!("{DATANODE_STAT_PREFIX}-0-").into_bytes()
387 }
388}
389
390impl From<DatanodeStatKey> for Vec<u8> {
391 fn from(value: DatanodeStatKey) -> Self {
392 format!("{}-0-{}", DATANODE_STAT_PREFIX, value.node_id).into_bytes()
394 }
395}
396
397impl FromStr for DatanodeStatKey {
398 type Err = error::Error;
399
400 fn from_str(key: &str) -> Result<Self> {
401 let caps = DATANODE_STAT_KEY_PATTERN
402 .captures(key)
403 .context(error::InvalidStatKeySnafu { key })?;
404
405 ensure!(caps.len() == 3, error::InvalidStatKeySnafu { key });
406 let node_id = caps[2].to_string();
407 let node_id: u64 = node_id.parse().context(error::ParseNumSnafu {
408 err_msg: format!("invalid node_id: {node_id}"),
409 })?;
410
411 Ok(Self { node_id })
412 }
413}
414
415impl TryFrom<Vec<u8>> for DatanodeStatKey {
416 type Error = error::Error;
417
418 fn try_from(bytes: Vec<u8>) -> Result<Self> {
419 String::from_utf8(bytes)
420 .context(error::FromUtf8Snafu {
421 name: "DatanodeStatKey",
422 })
423 .map(|x| x.parse())?
424 }
425}
426
427#[derive(Debug, Clone, Serialize, Deserialize)]
429#[serde(transparent)]
430pub struct DatanodeStatValue {
431 pub stats: Vec<Stat>,
432}
433
434impl DatanodeStatValue {
435 pub fn region_num(&self) -> Option<u64> {
437 self.stats.last().map(|x| x.region_num)
438 }
439
440 pub fn node_addr(&self) -> Option<String> {
442 self.stats.last().map(|x| x.addr.clone())
443 }
444}
445
446impl TryFrom<DatanodeStatValue> for Vec<u8> {
447 type Error = error::Error;
448
449 fn try_from(stats: DatanodeStatValue) -> Result<Self> {
450 Ok(serde_json::to_string(&stats)
451 .context(error::SerializeToJsonSnafu {
452 input: format!("{stats:?}"),
453 })?
454 .into_bytes())
455 }
456}
457
458impl FromStr for DatanodeStatValue {
459 type Err = error::Error;
460
461 fn from_str(value: &str) -> Result<Self> {
462 serde_json::from_str(value).context(error::DeserializeFromJsonSnafu { input: value })
463 }
464}
465
466impl TryFrom<Vec<u8>> for DatanodeStatValue {
467 type Error = error::Error;
468
469 fn try_from(value: Vec<u8>) -> Result<Self> {
470 String::from_utf8(value)
471 .context(error::FromUtf8Snafu {
472 name: "DatanodeStatValue",
473 })
474 .map(|x| x.parse())?
475 }
476}
477
478#[cfg(test)]
479mod tests {
480 use super::*;
481
482 #[test]
483 fn test_stat_key() {
484 let stat = Stat {
485 id: 101,
486 region_num: 10,
487 ..Default::default()
488 };
489
490 let stat_key = stat.stat_key();
491
492 assert_eq!(101, stat_key.node_id);
493 }
494
495 #[test]
496 fn test_stat_val_round_trip() {
497 let stat = Stat {
498 id: 101,
499 region_num: 100,
500 ..Default::default()
501 };
502
503 let stat_val = DatanodeStatValue { stats: vec![stat] };
504
505 let bytes: Vec<u8> = stat_val.try_into().unwrap();
506 let stat_val: DatanodeStatValue = bytes.try_into().unwrap();
507 let stats = stat_val.stats;
508
509 assert_eq!(1, stats.len());
510
511 let stat = stats.first().unwrap();
512 assert_eq!(101, stat.id);
513 assert_eq!(100, stat.region_num);
514 }
515
516 #[test]
517 fn test_get_addr_from_stat_val() {
518 let empty = DatanodeStatValue { stats: vec![] };
519 let addr = empty.node_addr();
520 assert!(addr.is_none());
521
522 let stat_val = DatanodeStatValue {
523 stats: vec![
524 Stat {
525 addr: "1".to_string(),
526 ..Default::default()
527 },
528 Stat {
529 addr: "2".to_string(),
530 ..Default::default()
531 },
532 Stat {
533 addr: "3".to_string(),
534 ..Default::default()
535 },
536 ],
537 };
538 let addr = stat_val.node_addr().unwrap();
539 assert_eq!("3", addr);
540 }
541
542 #[test]
543 fn test_get_region_num_from_stat_val() {
544 let empty = DatanodeStatValue { stats: vec![] };
545 let region_num = empty.region_num();
546 assert!(region_num.is_none());
547
548 let wrong = DatanodeStatValue {
549 stats: vec![Stat {
550 region_num: 0,
551 ..Default::default()
552 }],
553 };
554 let right = wrong.region_num();
555 assert_eq!(Some(0), right);
556
557 let stat_val = DatanodeStatValue {
558 stats: vec![
559 Stat {
560 region_num: 1,
561 ..Default::default()
562 },
563 Stat {
564 region_num: 0,
565 ..Default::default()
566 },
567 Stat {
568 region_num: 2,
569 ..Default::default()
570 },
571 ],
572 };
573 let region_num = stat_val.region_num().unwrap();
574 assert_eq!(2, region_num);
575 }
576}