below_model/
lib.rs

1// Copyright (c) Facebook, Inc. and its affiliates.
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
15#![deny(clippy::all)]
16
17use std::collections::BTreeMap;
18use std::collections::BTreeSet;
19use std::fmt;
20use std::marker::PhantomData;
21use std::str::FromStr;
22use std::time::Duration;
23use std::time::Instant;
24use std::time::SystemTime;
25
26use anyhow::anyhow;
27use anyhow::Context;
28use anyhow::Result;
29use common::open_source_shim;
30use enum_iterator::Sequence;
31use serde::Deserialize;
32use serde::Serialize;
33
34#[macro_use]
35pub mod collector;
36pub mod cgroup;
37pub mod collector_plugin;
38#[cfg(test)]
39mod common_field_ids;
40pub mod network;
41pub mod process;
42pub mod resctrl;
43pub mod sample;
44mod sample_model;
45pub mod system;
46pub mod tc_collector_plugin;
47pub mod tc_model;
48
49open_source_shim!(pub);
50
51pub use cgroup::*;
52pub use collector::*;
53pub use network::*;
54pub use process::*;
55pub use resctrl::*;
56pub use sample::*;
57pub use system::*;
58pub use tc_model::*;
59
60/// A wrapper for different field types used in Models. By this way we can query
61/// different fields in a single function without using Box.
62#[derive(Clone, Debug)]
63pub enum Field {
64    U32(u32),
65    U64(u64),
66    I32(i32),
67    I64(i64),
68    F32(f32),
69    F64(f64),
70    Str(String),
71    PidState(procfs::PidState),
72    VecU32(Vec<u32>),
73    StrSet(BTreeSet<String>),
74    StrU64Map(BTreeMap<String, u64>),
75    Cpuset(cgroupfs::Cpuset),
76    MemNodes(cgroupfs::MemNodes),
77    ResctrlCpuset(resctrlfs::Cpuset),
78    ResctrlGroupMode(resctrlfs::GroupMode),
79}
80
81impl From<Field> for u64 {
82    fn from(field: Field) -> u64 {
83        match field {
84            Field::U32(v) => v as u64,
85            Field::U64(v) => v,
86            _ => panic!("Operation for unsupported types"),
87        }
88    }
89}
90
91impl From<Field> for i64 {
92    fn from(field: Field) -> i64 {
93        match field {
94            Field::I32(v) => v as i64,
95            Field::I64(v) => v,
96            _ => panic!("Operation for unsupported types"),
97        }
98    }
99}
100
101impl From<Field> for f32 {
102    fn from(field: Field) -> f32 {
103        let result: f64 = field.into();
104        result as f32
105    }
106}
107
108impl From<Field> for f64 {
109    fn from(field: Field) -> f64 {
110        match field {
111            Field::U32(v) => v as f64,
112            Field::U64(v) => v as f64,
113            Field::I32(v) => v as f64,
114            Field::I64(v) => v as f64,
115            Field::F32(v) => v as f64,
116            Field::F64(v) => v,
117            _ => panic!("Operation for unsupported types"),
118        }
119    }
120}
121
122impl From<Field> for String {
123    fn from(field: Field) -> String {
124        match field {
125            Field::Str(v) => v,
126            _ => panic!("Operation for unsupported types"),
127        }
128    }
129}
130
131impl From<u32> for Field {
132    fn from(v: u32) -> Self {
133        Field::U32(v)
134    }
135}
136
137impl From<u64> for Field {
138    fn from(v: u64) -> Self {
139        Field::U64(v)
140    }
141}
142
143impl From<i32> for Field {
144    fn from(v: i32) -> Self {
145        Field::I32(v)
146    }
147}
148
149impl From<i64> for Field {
150    fn from(v: i64) -> Self {
151        Field::I64(v)
152    }
153}
154
155impl From<f32> for Field {
156    fn from(v: f32) -> Self {
157        Field::F32(v)
158    }
159}
160
161impl From<f64> for Field {
162    fn from(v: f64) -> Self {
163        Field::F64(v)
164    }
165}
166
167impl From<String> for Field {
168    fn from(v: String) -> Self {
169        Field::Str(v)
170    }
171}
172
173impl From<procfs::PidState> for Field {
174    fn from(v: procfs::PidState) -> Self {
175        Field::PidState(v)
176    }
177}
178
179impl From<Vec<u32>> for Field {
180    fn from(v: Vec<u32>) -> Self {
181        Field::VecU32(v)
182    }
183}
184
185impl From<BTreeSet<String>> for Field {
186    fn from(v: BTreeSet<String>) -> Self {
187        Field::StrSet(v)
188    }
189}
190
191impl From<BTreeMap<String, u64>> for Field {
192    fn from(v: BTreeMap<String, u64>) -> Self {
193        Field::StrU64Map(v)
194    }
195}
196
197impl From<cgroupfs::Cpuset> for Field {
198    fn from(v: cgroupfs::Cpuset) -> Self {
199        Field::Cpuset(v)
200    }
201}
202
203impl From<cgroupfs::MemNodes> for Field {
204    fn from(v: cgroupfs::MemNodes) -> Self {
205        Field::MemNodes(v)
206    }
207}
208
209impl From<resctrlfs::Cpuset> for Field {
210    fn from(v: resctrlfs::Cpuset) -> Self {
211        Field::ResctrlCpuset(v)
212    }
213}
214
215impl From<resctrlfs::GroupMode> for Field {
216    fn from(v: resctrlfs::GroupMode) -> Self {
217        Field::ResctrlGroupMode(v)
218    }
219}
220
221impl<T: Into<Field> + Clone> From<&T> for Field {
222    fn from(v: &T) -> Self {
223        v.clone().into()
224    }
225}
226
227impl std::ops::Add for Field {
228    type Output = Self;
229
230    fn add(self, other: Self) -> Self {
231        match (self, other) {
232            (Field::U32(s), Field::U32(o)) => (s + o).into(),
233            (Field::U64(s), Field::U64(o)) => (s + o).into(),
234            (Field::I32(s), Field::I32(o)) => (s + o).into(),
235            (Field::I64(s), Field::I64(o)) => (s + o).into(),
236            (Field::F32(s), Field::F32(o)) => (s + o).into(),
237            (Field::F64(s), Field::F64(o)) => (s + o).into(),
238            (Field::Str(s), Field::Str(o)) => (s + &o).into(),
239            _ => panic!("Operation for unsupported types"),
240        }
241    }
242}
243
244impl PartialEq for Field {
245    fn eq(&self, other: &Self) -> bool {
246        match (self, other) {
247            (Field::U32(s), Field::U32(o)) => s == o,
248            (Field::U64(s), Field::U64(o)) => s == o,
249            (Field::I32(s), Field::I32(o)) => s == o,
250            (Field::I64(s), Field::I64(o)) => s == o,
251            (Field::F32(s), Field::F32(o)) => s == o,
252            (Field::F64(s), Field::F64(o)) => s == o,
253            (Field::Str(s), Field::Str(o)) => s == o,
254            (Field::PidState(s), Field::PidState(o)) => s == o,
255            (Field::VecU32(s), Field::VecU32(o)) => s == o,
256            (Field::StrU64Map(s), Field::StrU64Map(o)) => s == o,
257            _ => false,
258        }
259    }
260}
261
262impl PartialOrd for Field {
263    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
264        match (self, other) {
265            (Field::U32(s), Field::U32(o)) => s.partial_cmp(o),
266            (Field::U64(s), Field::U64(o)) => s.partial_cmp(o),
267            (Field::I32(s), Field::I32(o)) => s.partial_cmp(o),
268            (Field::I64(s), Field::I64(o)) => s.partial_cmp(o),
269            (Field::F32(s), Field::F32(o)) => s.partial_cmp(o),
270            (Field::F64(s), Field::F64(o)) => s.partial_cmp(o),
271            (Field::Str(s), Field::Str(o)) => s.partial_cmp(o),
272            (Field::PidState(s), Field::PidState(o)) => s.partial_cmp(o),
273            (Field::VecU32(s), Field::VecU32(o)) => s.partial_cmp(o),
274            _ => None,
275        }
276    }
277}
278
279impl fmt::Display for Field {
280    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
281        match self {
282            Field::U32(v) => v.fmt(f),
283            Field::U64(v) => v.fmt(f),
284            Field::I32(v) => v.fmt(f),
285            Field::I64(v) => v.fmt(f),
286            Field::F32(v) => v.fmt(f),
287            Field::F64(v) => v.fmt(f),
288            Field::Str(v) => v.fmt(f),
289            Field::PidState(v) => v.fmt(f),
290            Field::VecU32(v) => f.write_fmt(format_args!("{:?}", v)),
291            Field::StrSet(v) => f.write_fmt(format_args!(
292                "{}",
293                v.iter()
294                    .cloned()
295                    .collect::<Vec<String>>()
296                    .as_slice()
297                    .join(" ")
298            )),
299            Field::StrU64Map(v) => f.write_fmt(format_args!(
300                "{}",
301                v.iter()
302                    .map(|(k, v)| format!("{}={}", k, v))
303                    .collect::<Vec<String>>()
304                    .as_slice()
305                    .join(", ")
306            )),
307            Field::Cpuset(v) => v.fmt(f),
308            Field::MemNodes(v) => v.fmt(f),
309            Field::ResctrlCpuset(v) => v.fmt(f),
310            Field::ResctrlGroupMode(v) => v.fmt(f),
311        }
312    }
313}
314
315/// Each Model is composed of Fields and optionally sub-Models. The Queriable
316/// trait let us query() a Model for a particular Field within the hierarchy
317/// with the given FieldId.
318pub trait Queriable {
319    type FieldId: FieldId<Queriable = Self>;
320    fn query(&self, field_id: &Self::FieldId) -> Option<Field>;
321}
322
323/// Marker trait to bind FieldId back to Queriable for type inference.
324pub trait FieldId: Sized {
325    type Queriable: Queriable<FieldId = Self> + ?Sized;
326}
327
328pub fn sort_queriables<T: Queriable>(queriables: &mut [&T], field_id: &T::FieldId, reverse: bool) {
329    queriables.sort_by(|lhs, rhs| {
330        let order = lhs
331            .query(field_id)
332            .partial_cmp(&rhs.query(field_id))
333            .unwrap_or(std::cmp::Ordering::Equal);
334        if reverse { order.reverse() } else { order }
335    });
336}
337
338/// Models containing sub-Models with its own type, similar to a node in a tree.
339/// Such Model has a depth value for illustrating the tree hierarchy.
340pub trait Recursive {
341    fn get_depth(&self) -> usize;
342}
343
344/// A `Model` that can be named
345pub trait Nameable {
346    /// The name of the model (for example "cgroup" or "system")
347    fn name() -> &'static str;
348}
349
350/// Type that contains sub-queriables of the same type, individually retrieveable
351/// by some index. It is itself a Queriable.
352pub trait QueriableContainer {
353    type Idx;
354    type SubqueryId: FieldId;
355    const IDX_PLACEHOLDER: &'static str = "<idx>.";
356    fn split(s: &str) -> Option<(&str, &str)> {
357        s.split_once('.')
358    }
359    fn get_item(&self, idx: &Self::Idx) -> Option<&<Self::SubqueryId as FieldId>::Queriable>;
360}
361
362impl<C: QueriableContainer> Queriable for C {
363    type FieldId = QueriableContainerFieldId<C>;
364    fn query(&self, field_id: &<C as Queriable>::FieldId) -> Option<Field> {
365        self.get_item(field_id.idx.as_ref()?)
366            .and_then(|sub| sub.query(&field_id.subquery_id.0))
367    }
368}
369
370#[derive(Clone, Debug, PartialEq)]
371pub struct QueriableContainerFieldId<C: QueriableContainer> {
372    /// None is only for listing variants and otherwise invalid.
373    /// If None, shows up as C::IDX_PLACEHOLDER
374    pub idx: Option<C::Idx>,
375    // Wraps inside a tuple so we can #[derive] traits without adding type constraints
376    pub subquery_id: (C::SubqueryId,),
377    phantom: PhantomData<C>,
378}
379
380impl<C: QueriableContainer> FieldId for QueriableContainerFieldId<C> {
381    type Queriable = C;
382}
383
384impl<C: QueriableContainer> QueriableContainerFieldId<C> {
385    pub fn new(idx: Option<C::Idx>, subquery_id: C::SubqueryId) -> Self {
386        Self {
387            idx,
388            subquery_id: (subquery_id,),
389            phantom: PhantomData,
390        }
391    }
392}
393
394impl<C: QueriableContainer> Sequence for QueriableContainerFieldId<C>
395where
396    C::SubqueryId: Sequence,
397{
398    const CARDINALITY: usize = C::SubqueryId::CARDINALITY;
399    fn next(&self) -> Option<Self> {
400        self.subquery_id.0.next().map(|s| Self::new(None, s))
401    }
402    fn previous(&self) -> Option<Self> {
403        self.subquery_id.0.previous().map(|s| Self::new(None, s))
404    }
405    fn first() -> Option<Self> {
406        C::SubqueryId::first().map(|s| Self::new(None, s))
407    }
408    fn last() -> Option<Self> {
409        C::SubqueryId::last().map(|s| Self::new(None, s))
410    }
411}
412
413impl<C: QueriableContainer> std::fmt::Display for QueriableContainerFieldId<C>
414where
415    C::Idx: ToString,
416    C::SubqueryId: ToString,
417{
418    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
419        match self.idx.as_ref() {
420            Some(idx) => write!(f, "{}.{}", idx.to_string(), self.subquery_id.0.to_string()),
421            None => write!(
422                f,
423                "{}{}",
424                C::IDX_PLACEHOLDER,
425                self.subquery_id.0.to_string()
426            ),
427        }
428    }
429}
430
431impl<C: QueriableContainer> FromStr for QueriableContainerFieldId<C>
432where
433    C::Idx: FromStr,
434    C::SubqueryId: FromStr,
435    <C::Idx as FromStr>::Err: Into<anyhow::Error>,
436    <C::SubqueryId as FromStr>::Err: Into<anyhow::Error>,
437{
438    type Err = anyhow::Error;
439    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
440        if let Some((idx_str, subquery_id_str)) = C::split(s) {
441            Ok(Self::new(
442                Some(C::Idx::from_str(idx_str).map_err(Into::into)?),
443                C::SubqueryId::from_str(subquery_id_str).map_err(Into::into)?,
444            ))
445        } else {
446            Err(anyhow!(
447                "Unable to find a variant of the given enum matching string `{}`.",
448                s,
449            ))
450        }
451    }
452}
453
454impl<Q: Queriable> QueriableContainer for Vec<Q> {
455    type Idx = usize;
456    type SubqueryId = Q::FieldId;
457    fn get_item(&self, idx: &usize) -> Option<&Q> {
458        self.get(*idx)
459    }
460}
461
462pub type VecFieldId<Q> = QueriableContainerFieldId<Vec<Q>>;
463
464impl<K: Ord, Q: Queriable> QueriableContainer for BTreeMap<K, Q> {
465    type Idx = K;
466    type SubqueryId = Q::FieldId;
467    const IDX_PLACEHOLDER: &'static str = "<key>.";
468    fn get_item(&self, idx: &K) -> Option<&Q> {
469        self.get(idx)
470    }
471}
472
473pub type BTreeMapFieldId<K, Q> = QueriableContainerFieldId<BTreeMap<K, Q>>;
474
475pub struct NetworkStats<'a> {
476    net: &'a procfs::NetStat,
477    ethtool: &'a Option<ethtool::EthtoolStats>,
478}
479
480#[derive(Clone, Serialize, Deserialize, below_derive::Queriable)]
481pub struct Model {
482    #[queriable(ignore)]
483    pub time_elapsed: Duration,
484    #[queriable(ignore)]
485    pub timestamp: SystemTime,
486    #[queriable(subquery)]
487    pub system: SystemModel,
488    #[queriable(subquery)]
489    pub cgroup: CgroupModel,
490    #[queriable(subquery)]
491    pub process: ProcessModel,
492    #[queriable(subquery)]
493    pub network: NetworkModel,
494    #[queriable(subquery)]
495    pub gpu: Option<GpuModel>,
496    #[queriable(subquery)]
497    pub resctrl: Option<ResctrlModel>,
498    #[queriable(subquery)]
499    pub tc: Option<TcModel>,
500}
501
502impl Model {
503    /// Construct a `Model` from a Sample and optionally, the last
504    /// `CumulativeSample` as well as the `Duration` since it was
505    /// collected.
506    pub fn new(timestamp: SystemTime, sample: &Sample, last: Option<(&Sample, Duration)>) -> Self {
507        Model {
508            time_elapsed: last.map(|(_, d)| d).unwrap_or_default(),
509            timestamp,
510            system: SystemModel::new(&sample.system, last.map(|(s, d)| (&s.system, d))),
511            cgroup: CgroupModel::new(
512                "<root>".to_string(),
513                String::new(),
514                0,
515                &sample.cgroup,
516                last.map(|(s, d)| (&s.cgroup, d)),
517            )
518            .aggr_top_level_val(),
519            process: ProcessModel::new(&sample.processes, last.map(|(s, d)| (&s.processes, d))),
520            network: {
521                let sample = NetworkStats {
522                    net: &sample.netstats,
523                    ethtool: &sample.ethtool,
524                };
525                let network_stats: NetworkStats;
526
527                let last = if let Some((s, d)) = last {
528                    network_stats = NetworkStats {
529                        net: &s.netstats,
530                        ethtool: &s.ethtool,
531                    };
532                    Some((&network_stats, d))
533                } else {
534                    None
535                };
536
537                NetworkModel::new(&sample, last)
538            },
539            gpu: sample.gpus.as_ref().map(|gpus| {
540                GpuModel::new(&gpus.gpu_map, {
541                    if let Some((s, d)) = last {
542                        s.gpus.as_ref().map(|g| (&g.gpu_map, d))
543                    } else {
544                        None
545                    }
546                })
547            }),
548            resctrl: sample.resctrl.as_ref().map(|r| {
549                ResctrlModel::new(
550                    r,
551                    if let Some((s, d)) = last {
552                        s.resctrl.as_ref().map(|r| (r, d))
553                    } else {
554                        None
555                    },
556                )
557            }),
558            tc: sample.tc.as_ref().map(|tc| {
559                TcModel::new(
560                    tc,
561                    if let Some((s, d)) = last {
562                        s.tc.as_ref().map(|tc| (tc, d))
563                    } else {
564                        None
565                    },
566                )
567            }),
568        }
569    }
570}
571
572/// Get a sample `Model`. There are no guarantees internal consistency of the
573/// model, neither are values in the model supposed to be realistic.
574pub fn get_sample_model() -> Model {
575    serde_json::from_str(sample_model::SAMPLE_MODEL_JSON)
576        .expect("Failed to deserialize sample model JSON")
577}
578
579#[cfg(test)]
580mod tests {
581    use super::*;
582
583    #[test]
584    fn test_model_field_ids() {
585        // Ensure COMMON_MODEL_FIELD_IDS is update to date.
586        let all_variants: BTreeSet<String> = enum_iterator::all::<ModelFieldId>()
587            .map(|v| v.to_string())
588            .collect();
589        let expected_field_ids: BTreeSet<String> = field_ids::MODEL_FIELD_IDS
590            .iter()
591            .map(|v| v.to_string())
592            .collect();
593
594        assert_eq!(
595            all_variants,
596            expected_field_ids,
597            "new fields: {:?}. missing fields: {:?}",
598            expected_field_ids
599                .difference(&all_variants)
600                .collect::<Vec<_>>(),
601            all_variants
602                .difference(&expected_field_ids)
603                .collect::<Vec<_>>()
604        );
605    }
606
607    #[test]
608    fn test_deserialize_sample_model_json() {
609        get_sample_model();
610    }
611
612    #[::below_derive::queriable_derives]
613    pub struct TestModel {
614        pub msg: String,
615    }
616
617    #[test]
618    fn test_vec_field_id() {
619        let query_str = "1.msg";
620        let query = <VecFieldId<TestModel>>::from_str(query_str).expect("bad query str");
621        assert_eq!(query, VecFieldId::new(Some(1), TestModelFieldId::Msg),);
622        assert_eq!(query.to_string(), query_str);
623    }
624
625    #[test]
626    fn test_query_vec() {
627        let data = vec![
628            TestModel {
629                msg: "hello".to_owned(),
630            },
631            TestModel {
632                msg: "world".to_owned(),
633            },
634        ];
635        assert_eq!(
636            data.query(&VecFieldId::new(Some(1), TestModelFieldId::Msg,)),
637            Some(Field::Str("world".to_owned()))
638        );
639    }
640
641    #[test]
642    fn test_btreemap_field_id() {
643        let query_str = "hello.msg";
644        let query =
645            <BTreeMapFieldId<String, TestModel>>::from_str(query_str).expect("bad query str");
646        assert_eq!(
647            query,
648            BTreeMapFieldId::new(Some("hello".to_owned()), TestModelFieldId::Msg)
649        );
650        assert_eq!(query.to_string(), query_str);
651    }
652
653    #[test]
654    fn test_query_btreemap() {
655        let mut data = <BTreeMap<String, TestModel>>::new();
656        data.insert(
657            "hello".to_owned(),
658            TestModel {
659                msg: "world".to_owned(),
660            },
661        );
662        data.insert(
663            "foo".to_owned(),
664            TestModel {
665                msg: "bar".to_owned(),
666            },
667        );
668        assert_eq!(
669            data.query(&BTreeMapFieldId::new(
670                Some("hello".to_owned()),
671                TestModelFieldId::Msg,
672            )),
673            Some(Field::Str("world".to_owned()))
674        );
675    }
676
677    #[test]
678    fn test_query_models() {
679        let model = get_sample_model();
680        for (field_id, expected) in &[
681            (
682                "system.hostname",
683                Some(Field::Str("hostname.example.com".to_owned())),
684            ),
685            (
686                "cgroup.path:/init.scope/.cpu.usage_pct",
687                Some(Field::F64(0.01)),
688            ),
689            (
690                "network.interfaces.eth0.interface",
691                Some(Field::Str("eth0".to_owned())),
692            ),
693            (
694                "process.processes.1.comm",
695                Some(Field::Str("systemd".to_owned())),
696            ),
697        ] {
698            assert_eq!(
699                &model.query(
700                    &ModelFieldId::from_str(field_id)
701                        .map_err(|e| format!("Failed to parse field id {}: {:?}", field_id, e))
702                        .unwrap()
703                ),
704                expected
705            );
706        }
707    }
708}
below_model/lib.rs

below_model/
lib.rs
