common_meta/
datanode.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::collections::{HashMap, HashSet};
16use std::str::FromStr;
17
18use api::v1::meta::{DatanodeWorkloads, HeartbeatRequest, RequestHeader};
19use common_time::{Timestamp, 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/// The private table name for persisted region statistics.
35pub const REGION_STATS_HISTORY_TABLE_NAME: &str = "region_statistics_history";
36
37lazy_static! {
38    pub(crate) static ref DATANODE_LEASE_KEY_PATTERN: Regex =
39        Regex::new("^__meta_datanode_lease-([0-9]+)-([0-9]+)$").unwrap();
40    static ref DATANODE_STAT_KEY_PATTERN: Regex =
41        Regex::new(&format!("^{DATANODE_STAT_PREFIX}-([0-9]+)-([0-9]+)$")).unwrap();
42    static ref INACTIVE_REGION_KEY_PATTERN: Regex =
43        Regex::new("^__meta_inactive_region-([0-9]+)-([0-9]+)-([0-9]+)$").unwrap();
44}
45
46/// The key of the datanode stat in the storage.
47///
48/// The format is `__meta_datanode_stat-0-{node_id}`.
49#[derive(Debug, Clone, Default, Serialize, Deserialize)]
50pub struct Stat {
51    pub timestamp_millis: i64,
52    // The datanode Id.
53    pub id: u64,
54    // The datanode address.
55    pub addr: String,
56    /// The read capacity units during this period
57    pub rcus: i64,
58    /// The write capacity units during this period
59    pub wcus: i64,
60    /// How many regions on this node
61    pub region_num: u64,
62    /// The region stats of the datanode.
63    pub region_stats: Vec<RegionStat>,
64    /// The topic stats of the datanode.
65    pub topic_stats: Vec<TopicStat>,
66    // The node epoch is used to check whether the node has restarted or redeployed.
67    pub node_epoch: u64,
68    /// The datanode workloads.
69    pub datanode_workloads: DatanodeWorkloads,
70    /// The GC statistics of the datanode.
71    pub gc_stat: Option<GcStat>,
72}
73
74/// The statistics of a region.
75#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
76pub struct RegionStat {
77    /// The region_id.
78    pub id: RegionId,
79    /// The read capacity units during this period
80    pub rcus: i64,
81    /// The write capacity units during this period
82    pub wcus: i64,
83    /// Approximate disk bytes of this region, including sst, index, manifest and wal
84    pub approximate_bytes: u64,
85    /// The engine name.
86    pub engine: String,
87    /// The region role.
88    pub role: RegionRole,
89    /// The number of rows
90    pub num_rows: u64,
91    /// The size of the memtable in bytes.
92    pub memtable_size: u64,
93    /// The size of the manifest in bytes.
94    pub manifest_size: u64,
95    /// The size of the SST data files in bytes.
96    pub sst_size: u64,
97    /// The num of the SST data files.
98    pub sst_num: u64,
99    /// The size of the SST index files in bytes.
100    pub index_size: u64,
101    /// The manifest infoof the region.
102    pub region_manifest: RegionManifestInfo,
103    /// The total bytes written of the region since region opened.
104    pub written_bytes: u64,
105    /// Total ingested Parquet bytes since the region opened.
106    /// Currently counts only committed flush output.
107    ///
108    /// Excludes index files, compaction and imported SSTs. Resets on region reopen.
109    /// Accepts the legacy `flushed_parquet_bytes` key; upgrade metasrv before datanodes.
110    #[serde(default, alias = "flushed_parquet_bytes")]
111    pub ingested_parquet_bytes: u64,
112    /// The total query CPU time of the region since region opened.
113    ///
114    /// Unit: nanoseconds.
115    #[serde(default)]
116    pub query_cpu_time: u64,
117    /// The total scanned bytes of the region since region opened.
118    #[serde(default)]
119    pub query_scanned_bytes: u64,
120    /// The latest entry id of topic used by data.
121    /// **Only used by remote WAL prune.**
122    pub data_topic_latest_entry_id: u64,
123    /// The latest entry id of topic used by metadata.
124    /// **Only used by remote WAL prune.**
125    /// In mito engine, this is the same as `data_topic_latest_entry_id`.
126    pub metadata_topic_latest_entry_id: u64,
127    /// The earliest timestamp of the region's current data, if it holds any.
128    #[serde(default)]
129    pub min_timestamp: Option<Timestamp>,
130    /// The latest timestamp of the region's current data, if it holds any.
131    #[serde(default)]
132    pub max_timestamp: Option<Timestamp>,
133}
134
135#[derive(Debug, Clone, Serialize, Deserialize)]
136pub struct TopicStat {
137    /// The topic name.
138    pub topic: String,
139    /// The latest entry id of the topic.
140    pub latest_entry_id: u64,
141    /// The total size in bytes of records appended to the topic.
142    pub record_size: u64,
143    /// The total number of records appended to the topic.
144    pub record_num: u64,
145}
146
147/// Trait for reporting statistics about topics.
148pub trait TopicStatsReporter: Send + Sync {
149    /// Returns a list of topic statistics that can be reported.
150    fn reportable_topics(&mut self) -> Vec<TopicStat>;
151}
152
153#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
154pub enum RegionManifestInfo {
155    Mito {
156        manifest_version: u64,
157        flushed_entry_id: u64,
158        /// Number of files removed in the manifest's `removed_files` field.
159        file_removed_cnt: u64,
160    },
161    Metric {
162        data_manifest_version: u64,
163        data_flushed_entry_id: u64,
164        metadata_manifest_version: u64,
165        metadata_flushed_entry_id: u64,
166    },
167}
168
169impl Stat {
170    #[inline]
171    pub fn is_empty(&self) -> bool {
172        self.region_stats.is_empty()
173    }
174
175    pub fn stat_key(&self) -> DatanodeStatKey {
176        DatanodeStatKey { node_id: self.id }
177    }
178
179    /// Returns a tuple array containing [RegionId] and [RegionRole].
180    pub fn regions(&self) -> Vec<(RegionId, RegionRole)> {
181        self.region_stats.iter().map(|s| (s.id, s.role)).collect()
182    }
183
184    /// Returns all table ids in the region stats.
185    pub fn table_ids(&self) -> HashSet<TableId> {
186        self.region_stats.iter().map(|s| s.id.table_id()).collect()
187    }
188
189    /// Retains the active region stats and updates the rcus, wcus, and region_num.
190    pub fn retain_active_region_stats(&mut self, inactive_region_ids: &HashSet<RegionId>) {
191        if inactive_region_ids.is_empty() {
192            return;
193        }
194
195        self.region_stats
196            .retain(|r| !inactive_region_ids.contains(&r.id));
197        self.rcus = self.region_stats.iter().map(|s| s.rcus).sum();
198        self.wcus = self.region_stats.iter().map(|s| s.wcus).sum();
199        self.region_num = self.region_stats.len() as u64;
200    }
201
202    pub fn memory_size(&self) -> usize {
203        // timestamp_millis, rcus, wcus
204        std::mem::size_of::<i64>() * 3 +
205        // id, region_num, node_epoch
206        std::mem::size_of::<u64>() * 3 +
207        // addr
208        std::mem::size_of::<String>() + self.addr.capacity() +
209        // region_stats
210        self.region_stats.iter().map(|s| s.memory_size()).sum::<usize>()
211    }
212}
213
214impl RegionStat {
215    pub fn memory_size(&self) -> usize {
216        // role
217        std::mem::size_of::<RegionRole>() +
218        // id
219        std::mem::size_of::<RegionId>() +
220        // rcus, wcus, approximate_bytes, num_rows
221        std::mem::size_of::<i64>() * 4 +
222        // memtable_size, manifest_size, sst_size, sst_num, index_size
223        std::mem::size_of::<u64>() * 5 +
224        // engine
225        std::mem::size_of::<String>() + self.engine.capacity() +
226        // region_manifest
227        self.region_manifest.memory_size()
228    }
229}
230
231impl RegionManifestInfo {
232    pub fn memory_size(&self) -> usize {
233        match self {
234            RegionManifestInfo::Mito { .. } => std::mem::size_of::<u64>() * 2,
235            RegionManifestInfo::Metric { .. } => std::mem::size_of::<u64>() * 4,
236        }
237    }
238}
239
240impl TryFrom<&HeartbeatRequest> for Stat {
241    type Error = Option<RequestHeader>;
242
243    fn try_from(value: &HeartbeatRequest) -> std::result::Result<Self, Self::Error> {
244        let HeartbeatRequest {
245            header,
246            peer,
247            region_stats,
248            node_epoch,
249            node_workloads,
250            topic_stats,
251            extensions,
252            ..
253        } = value;
254
255        match (header, peer) {
256            (Some(header), Some(peer)) => {
257                let region_stats = region_stats
258                    .iter()
259                    .map(RegionStat::from)
260                    .collect::<Vec<_>>();
261                let topic_stats = topic_stats.iter().map(TopicStat::from).collect::<Vec<_>>();
262
263                let datanode_workloads = get_datanode_workloads(node_workloads.as_ref());
264
265                let gc_stat = GcStat::from_extensions(extensions).map_err(|err| {
266                    common_telemetry::error!(
267                        "Failed to deserialize GcStat from extensions: {}",
268                        err
269                    );
270                    header.clone()
271                })?;
272                Ok(Self {
273                    timestamp_millis: time_util::current_time_millis(),
274                    // datanode id
275                    id: peer.id,
276                    // datanode address
277                    addr: peer.addr.clone(),
278                    rcus: region_stats.iter().map(|s| s.rcus).sum(),
279                    wcus: region_stats.iter().map(|s| s.wcus).sum(),
280                    region_num: region_stats.len() as u64,
281                    region_stats,
282                    topic_stats,
283                    node_epoch: *node_epoch,
284                    datanode_workloads,
285                    gc_stat,
286                })
287            }
288            (header, _) => Err(header.clone()),
289        }
290    }
291}
292
293impl From<store_api::region_engine::RegionManifestInfo> for RegionManifestInfo {
294    fn from(value: store_api::region_engine::RegionManifestInfo) -> Self {
295        match value {
296            store_api::region_engine::RegionManifestInfo::Mito {
297                manifest_version,
298                flushed_entry_id,
299                file_removed_cnt,
300            } => RegionManifestInfo::Mito {
301                manifest_version,
302                flushed_entry_id,
303                file_removed_cnt,
304            },
305            store_api::region_engine::RegionManifestInfo::Metric {
306                data_manifest_version,
307                data_flushed_entry_id,
308                metadata_manifest_version,
309                metadata_flushed_entry_id,
310            } => RegionManifestInfo::Metric {
311                data_manifest_version,
312                data_flushed_entry_id,
313                metadata_manifest_version,
314                metadata_flushed_entry_id,
315            },
316        }
317    }
318}
319
320impl From<&api::v1::meta::RegionStat> for RegionStat {
321    fn from(value: &api::v1::meta::RegionStat) -> Self {
322        let region_stat = value
323            .extensions
324            .get(REGION_STATISTIC_KEY)
325            .and_then(|value| RegionStatistic::deserialize_from_slice(value))
326            .unwrap_or_default();
327
328        Self {
329            id: RegionId::from_u64(value.region_id),
330            rcus: value.rcus,
331            wcus: value.wcus,
332            approximate_bytes: value.approximate_bytes as u64,
333            engine: value.engine.clone(),
334            role: RegionRole::from(value.role()),
335            num_rows: region_stat.num_rows,
336            memtable_size: region_stat.memtable_size,
337            manifest_size: region_stat.manifest_size,
338            sst_size: region_stat.sst_size,
339            sst_num: region_stat.sst_num,
340            index_size: region_stat.index_size,
341            region_manifest: region_stat.manifest.into(),
342            written_bytes: region_stat.written_bytes,
343            ingested_parquet_bytes: region_stat.ingested_parquet_bytes,
344            query_cpu_time: region_stat.query_cpu_time,
345            query_scanned_bytes: region_stat.query_scanned_bytes,
346            data_topic_latest_entry_id: region_stat.data_topic_latest_entry_id,
347            metadata_topic_latest_entry_id: region_stat.metadata_topic_latest_entry_id,
348            min_timestamp: region_stat.min_timestamp,
349            max_timestamp: region_stat.max_timestamp,
350        }
351    }
352}
353
354impl From<&api::v1::meta::TopicStat> for TopicStat {
355    fn from(value: &api::v1::meta::TopicStat) -> Self {
356        Self {
357            topic: value.topic_name.clone(),
358            latest_entry_id: value.latest_entry_id,
359            record_size: value.record_size,
360            record_num: value.record_num,
361        }
362    }
363}
364
365#[derive(Debug, Clone, Serialize, Deserialize, Default)]
366pub struct GcStat {
367    /// Number of GC tasks currently running on the datanode.
368    pub running_gc_tasks: u32,
369    /// The maximum number of concurrent GC tasks the datanode can handle.
370    pub gc_concurrency: u32,
371}
372
373impl GcStat {
374    pub const GC_STAT_KEY: &str = "__gc_stat";
375
376    pub fn new(running_gc_tasks: u32, gc_concurrency: u32) -> Self {
377        Self {
378            running_gc_tasks,
379            gc_concurrency,
380        }
381    }
382
383    pub fn into_extensions(&self, extensions: &mut std::collections::HashMap<String, Vec<u8>>) {
384        let bytes = serde_json::to_vec(self).unwrap_or_default();
385        extensions.insert(Self::GC_STAT_KEY.to_string(), bytes);
386    }
387
388    pub fn from_extensions(
389        extensions: &std::collections::HashMap<String, Vec<u8>>,
390    ) -> Result<Option<Self>> {
391        extensions
392            .get(Self::GC_STAT_KEY)
393            .map(|bytes| {
394                serde_json::from_slice(bytes).with_context(|_| DeserializeFromJsonSnafu {
395                    input: String::from_utf8_lossy(bytes).to_string(),
396                })
397            })
398            .transpose()
399    }
400}
401
402/// Environment variables reported by a node in heartbeat messages.
403#[derive(Debug, Clone, Serialize, Deserialize, Default)]
404pub struct EnvVars {
405    pub vars: HashMap<String, String>,
406}
407
408impl EnvVars {
409    pub const ENV_VARS_KEY: &str = "__env_vars";
410
411    pub fn new(vars: HashMap<String, String>) -> Self {
412        Self { vars }
413    }
414
415    /// Read the configured env var keys from the environment and build an EnvVars.
416    pub fn from_config(keys: &[String]) -> Self {
417        let vars = keys
418            .iter()
419            .filter_map(|key| std::env::var(key).ok().map(|value| (key.clone(), value)))
420            .collect();
421        Self { vars }
422    }
423
424    pub fn into_extensions(&self, extensions: &mut HashMap<String, Vec<u8>>) {
425        if self.vars.is_empty() {
426            return;
427        }
428        let bytes = serde_json::to_vec(self).unwrap_or_default();
429        extensions.insert(Self::ENV_VARS_KEY.to_string(), bytes);
430    }
431
432    pub fn from_extensions(extensions: &HashMap<String, Vec<u8>>) -> Result<Option<Self>> {
433        extensions
434            .get(Self::ENV_VARS_KEY)
435            .map(|bytes| {
436                serde_json::from_slice(bytes).with_context(|_| DeserializeFromJsonSnafu {
437                    input: String::from_utf8_lossy(bytes).to_string(),
438                })
439            })
440            .transpose()
441    }
442}
443
444/// The key of the datanode stat in the memory store.
445///
446/// The format is `__meta_datanode_stat-0-{node_id}`.
447#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
448pub struct DatanodeStatKey {
449    pub node_id: u64,
450}
451
452impl DatanodeStatKey {
453    /// The key prefix.
454    pub fn prefix_key() -> Vec<u8> {
455        // todo(hl): remove cluster id in prefix
456        format!("{DATANODE_STAT_PREFIX}-0-").into_bytes()
457    }
458}
459
460impl From<DatanodeStatKey> for Vec<u8> {
461    fn from(value: DatanodeStatKey) -> Self {
462        // todo(hl): remove cluster id in prefix
463        format!("{}-0-{}", DATANODE_STAT_PREFIX, value.node_id).into_bytes()
464    }
465}
466
467impl FromStr for DatanodeStatKey {
468    type Err = error::Error;
469
470    fn from_str(key: &str) -> Result<Self> {
471        let caps = DATANODE_STAT_KEY_PATTERN
472            .captures(key)
473            .context(error::InvalidStatKeySnafu { key })?;
474
475        ensure!(caps.len() == 3, error::InvalidStatKeySnafu { key });
476        let node_id = caps[2].to_string();
477        let node_id: u64 = node_id.parse().context(error::ParseNumSnafu {
478            err_msg: format!("invalid node_id: {node_id}"),
479        })?;
480
481        Ok(Self { node_id })
482    }
483}
484
485impl TryFrom<Vec<u8>> for DatanodeStatKey {
486    type Error = error::Error;
487
488    fn try_from(bytes: Vec<u8>) -> Result<Self> {
489        String::from_utf8(bytes)
490            .context(error::FromUtf8Snafu {
491                name: "DatanodeStatKey",
492            })
493            .map(|x| x.parse())?
494    }
495}
496
497/// The value of the datanode stat in the memory store.
498#[derive(Debug, Clone, Serialize, Deserialize)]
499#[serde(transparent)]
500pub struct DatanodeStatValue {
501    pub stats: Vec<Stat>,
502}
503
504impl DatanodeStatValue {
505    /// Get the latest number of regions.
506    pub fn region_num(&self) -> Option<u64> {
507        self.stats.last().map(|x| x.region_num)
508    }
509
510    /// Get the latest node addr.
511    pub fn node_addr(&self) -> Option<String> {
512        self.stats.last().map(|x| x.addr.clone())
513    }
514}
515
516impl TryFrom<DatanodeStatValue> for Vec<u8> {
517    type Error = error::Error;
518
519    fn try_from(stats: DatanodeStatValue) -> Result<Self> {
520        Ok(serde_json::to_string(&stats)
521            .context(error::SerializeToJsonSnafu {
522                input: format!("{stats:?}"),
523            })?
524            .into_bytes())
525    }
526}
527
528impl FromStr for DatanodeStatValue {
529    type Err = error::Error;
530
531    fn from_str(value: &str) -> Result<Self> {
532        serde_json::from_str(value).context(error::DeserializeFromJsonSnafu { input: value })
533    }
534}
535
536impl TryFrom<Vec<u8>> for DatanodeStatValue {
537    type Error = error::Error;
538
539    fn try_from(value: Vec<u8>) -> Result<Self> {
540        String::from_utf8(value)
541            .context(error::FromUtf8Snafu {
542                name: "DatanodeStatValue",
543            })
544            .map(|x| x.parse())?
545    }
546}
547
548#[cfg(test)]
549mod tests {
550    use super::*;
551
552    #[test]
553    fn test_stat_key() {
554        let stat = Stat {
555            id: 101,
556            region_num: 10,
557            ..Default::default()
558        };
559
560        let stat_key = stat.stat_key();
561
562        assert_eq!(101, stat_key.node_id);
563    }
564
565    #[test]
566    fn test_stat_val_round_trip() {
567        let stat = Stat {
568            id: 101,
569            region_num: 100,
570            ..Default::default()
571        };
572
573        let stat_val = DatanodeStatValue { stats: vec![stat] };
574
575        let bytes: Vec<u8> = stat_val.try_into().unwrap();
576        let stat_val: DatanodeStatValue = bytes.try_into().unwrap();
577        let stats = stat_val.stats;
578
579        assert_eq!(1, stats.len());
580
581        let stat = stats.first().unwrap();
582        assert_eq!(101, stat.id);
583        assert_eq!(100, stat.region_num);
584    }
585
586    #[test]
587    fn test_stat_val_deserializes_without_optional_counters() {
588        let stat = Stat {
589            region_stats: vec![RegionStat {
590                id: RegionId::new(1024, 1),
591                rcus: 0,
592                wcus: 0,
593                approximate_bytes: 0,
594                engine: "mito".to_string(),
595                role: RegionRole::Leader,
596                num_rows: 0,
597                memtable_size: 0,
598                manifest_size: 0,
599                sst_size: 0,
600                sst_num: 0,
601                index_size: 0,
602                region_manifest: RegionManifestInfo::Mito {
603                    manifest_version: 0,
604                    flushed_entry_id: 0,
605                    file_removed_cnt: 0,
606                },
607                written_bytes: 0,
608                ingested_parquet_bytes: 4096,
609                query_cpu_time: 10,
610                query_scanned_bytes: 20,
611                data_topic_latest_entry_id: 0,
612                metadata_topic_latest_entry_id: 0,
613                min_timestamp: None,
614                max_timestamp: None,
615            }],
616            ..Default::default()
617        };
618        let stat_val = DatanodeStatValue { stats: vec![stat] };
619        let mut value = serde_json::to_value(stat_val).unwrap();
620        let region_stat = value[0]["region_stats"][0].as_object_mut().unwrap();
621        region_stat.remove("ingested_parquet_bytes");
622        region_stat.remove("query_cpu_time");
623        region_stat.remove("query_scanned_bytes");
624
625        let stat_val: DatanodeStatValue = serde_json::from_value(value).unwrap();
626        let region_stat = &stat_val.stats[0].region_stats[0];
627
628        assert_eq!(region_stat.ingested_parquet_bytes, 0);
629        assert_eq!(region_stat.query_cpu_time, 0);
630        assert_eq!(region_stat.query_scanned_bytes, 0);
631    }
632
633    #[test]
634    fn test_get_addr_from_stat_val() {
635        let empty = DatanodeStatValue { stats: vec![] };
636        let addr = empty.node_addr();
637        assert!(addr.is_none());
638
639        let stat_val = DatanodeStatValue {
640            stats: vec![
641                Stat {
642                    addr: "1".to_string(),
643                    ..Default::default()
644                },
645                Stat {
646                    addr: "2".to_string(),
647                    ..Default::default()
648                },
649                Stat {
650                    addr: "3".to_string(),
651                    ..Default::default()
652                },
653            ],
654        };
655        let addr = stat_val.node_addr().unwrap();
656        assert_eq!("3", addr);
657    }
658
659    #[test]
660    fn test_get_region_num_from_stat_val() {
661        let empty = DatanodeStatValue { stats: vec![] };
662        let region_num = empty.region_num();
663        assert!(region_num.is_none());
664
665        let wrong = DatanodeStatValue {
666            stats: vec![Stat {
667                region_num: 0,
668                ..Default::default()
669            }],
670        };
671        let right = wrong.region_num();
672        assert_eq!(Some(0), right);
673
674        let stat_val = DatanodeStatValue {
675            stats: vec![
676                Stat {
677                    region_num: 1,
678                    ..Default::default()
679                },
680                Stat {
681                    region_num: 0,
682                    ..Default::default()
683                },
684                Stat {
685                    region_num: 2,
686                    ..Default::default()
687                },
688            ],
689        };
690        let region_num = stat_val.region_num().unwrap();
691        assert_eq!(2, region_num);
692    }
693
694    #[test]
695    fn test_region_stat_from_heartbeat_preserves_staging_leader_role() {
696        let request = HeartbeatRequest {
697            header: Some(RequestHeader::default()),
698            peer: Some(api::v1::meta::Peer {
699                id: 1,
700                addr: "127.0.0.1:3001".to_string(),
701            }),
702            region_stats: vec![api::v1::meta::RegionStat {
703                region_id: RegionId::new(1024, 1).as_u64(),
704                engine: "mito".to_string(),
705                role: api::v1::meta::RegionRole::StagingLeader.into(),
706                ..Default::default()
707            }],
708            ..Default::default()
709        };
710
711        let stat = Stat::try_from(&request).unwrap();
712
713        assert_eq!(stat.region_stats.len(), 1);
714        assert_eq!(stat.region_stats[0].role, RegionRole::StagingLeader);
715    }
716
717    #[test]
718    fn test_ingested_parquet_bytes_from_heartbeat() {
719        for bytes in [0, 4096, u64::MAX] {
720            let statistic = RegionStatistic {
721                ingested_parquet_bytes: bytes,
722                ..Default::default()
723            };
724            let mut region = api::v1::meta::RegionStat::default();
725            region.extensions.insert(
726                REGION_STATISTIC_KEY.to_string(),
727                serde_json::to_vec(&statistic).unwrap(),
728            );
729            let request = HeartbeatRequest {
730                header: Some(RequestHeader::default()),
731                peer: Some(api::v1::meta::Peer {
732                    id: 1,
733                    addr: "127.0.0.1:3001".to_string(),
734                }),
735                region_stats: vec![region.clone()],
736                ..Default::default()
737            };
738            let stat = Stat::try_from(&request).unwrap();
739            let stat = &stat.region_stats[0];
740            assert_eq!(stat.ingested_parquet_bytes, bytes);
741            let encoded = serde_json::to_vec(stat).unwrap();
742            let wire: serde_json::Value = serde_json::from_slice(&encoded).unwrap();
743            assert_eq!(wire["ingested_parquet_bytes"], bytes);
744            assert!(wire.get("flushed_parquet_bytes").is_none());
745            let mut old_wire = wire.clone();
746            old_wire
747                .as_object_mut()
748                .unwrap()
749                .remove("ingested_parquet_bytes");
750            old_wire["flushed_parquet_bytes"] = bytes.into();
751            assert_eq!(
752                &serde_json::from_value::<RegionStat>(old_wire).unwrap(),
753                stat
754            );
755            let decoded: RegionStat = serde_json::from_slice(&encoded).unwrap();
756            assert_eq!(&decoded, stat);
757
758            let mut legacy = serde_json::to_value(&statistic).unwrap();
759            legacy
760                .as_object_mut()
761                .unwrap()
762                .remove("ingested_parquet_bytes");
763            legacy["flushed_parquet_bytes"] = bytes.into();
764            region.extensions.insert(
765                REGION_STATISTIC_KEY.to_string(),
766                serde_json::to_vec(&legacy).unwrap(),
767            );
768            assert_eq!(RegionStat::from(&region).ingested_parquet_bytes, bytes);
769            legacy
770                .as_object_mut()
771                .unwrap()
772                .remove("flushed_parquet_bytes");
773            region.extensions.insert(
774                REGION_STATISTIC_KEY.to_string(),
775                serde_json::to_vec(&legacy).unwrap(),
776            );
777            assert_eq!(RegionStat::from(&region).ingested_parquet_bytes, 0);
778        }
779        assert_eq!(
780            RegionStat::from(&api::v1::meta::RegionStat::default()).ingested_parquet_bytes,
781            0
782        );
783    }
784
785    #[test]
786    fn test_env_vars_round_trip() {
787        let mut vars = HashMap::new();
788        vars.insert("AZ".to_string(), "us-east-1a".to_string());
789        vars.insert("REGION".to_string(), "us-east-1".to_string());
790        let env_vars = EnvVars::new(vars);
791
792        let mut extensions = HashMap::new();
793        env_vars.into_extensions(&mut extensions);
794
795        let extracted = EnvVars::from_extensions(&extensions).unwrap().unwrap();
796        assert_eq!(extracted.vars.get("AZ").unwrap(), "us-east-1a");
797        assert_eq!(extracted.vars.get("REGION").unwrap(), "us-east-1");
798    }
799
800    #[test]
801    fn test_env_vars_empty_not_written() {
802        let env_vars = EnvVars::default();
803        let mut extensions = HashMap::new();
804        env_vars.into_extensions(&mut extensions);
805        assert!(extensions.is_empty());
806    }
807
808    #[test]
809    fn test_env_vars_from_extensions_missing() {
810        let extensions = HashMap::new();
811        let result = EnvVars::from_extensions(&extensions).unwrap();
812        assert!(result.is_none());
813    }
814}
common_meta/datanode.rs

common_meta/
datanode.rs
