api-tools 0.7.0

An API tools library for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196

//! Prometheus' metrics layer

use axum::body::Body;
use axum::extract::MatchedPath;
use axum::http::Request;
use axum::response::Response;
use bytesize::ByteSize;
use futures::future::BoxFuture;
use metrics::{counter, gauge, histogram};
use std::fmt;
use std::path::PathBuf;
use std::task::{Context, Poll};
use std::time::Instant;
use sysinfo::{Disks, System};
use tower::{Layer, Service};

/// Prometheus metrics layer for Axum
#[derive(Clone)]
pub struct PrometheusLayer {
    /// Service name
    pub service_name: String,

    /// Disk mount points to monitor
    pub disk_mount_points: Vec<PathBuf>,
}

impl<S> Layer<S> for PrometheusLayer {
    type Service = PrometheusMiddleware<S>;

    fn layer(&self, inner: S) -> Self::Service {
        PrometheusMiddleware {
            inner,
            service_name: self.service_name.clone(),
            disk_mount_points: self.disk_mount_points.clone(),
        }
    }
}

#[derive(Clone)]
pub struct PrometheusMiddleware<S> {
    inner: S,
    service_name: String,
    disk_mount_points: Vec<PathBuf>,
}

impl<S> Service<Request<Body>> for PrometheusMiddleware<S>
where
    S: Service<Request<Body>, Response = Response> + Send + 'static,
    S::Future: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    // `BoxFuture` is a type alias for `Pin<Box<dyn Future + Send + 'a>>`
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, request: Request<Body>) -> Self::Future {
        let path = if let Some(matched_path) = request.extensions().get::<MatchedPath>() {
            matched_path.as_str().to_owned()
        } else {
            request.uri().path().to_owned()
        };
        let method = request.method().to_string();
        let service_name = self.service_name.clone();
        let disk_mount_points = self.disk_mount_points.clone();

        let start = Instant::now();
        let future = self.inner.call(request);
        Box::pin(async move {
            let response = future.await?;

            // Exclude metrics endpoint
            if path != "/metrics" {
                let latency = start.elapsed().as_secs_f64();
                let status = response.status().as_u16().to_string();
                let labels = [
                    ("method", method),
                    ("path", path),
                    ("service", service_name.clone()),
                    ("status", status),
                ];

                counter!("http_requests_total", &labels).increment(1);
                histogram!("http_requests_duration_seconds", &labels).record(latency);
            }

            // System metrics
            let system_metrics = SystemMetrics::new(&disk_mount_points).await;
            system_metrics.add_metrics(service_name);

            Ok(response)
        })
    }
}

#[derive(Debug, Clone)]
struct SystemMetrics {
    /// Average CPU usage in percent
    cpu_usage: f32,

    /// Total memory in bytes
    total_memory: u64,

    /// Used memory in bytes
    used_memory: u64,

    /// Total swap space in bytes
    total_swap: u64,

    /// Used swap space in bytes
    used_swap: u64,

    /// Total disk space in bytes for a specified mount point
    total_disks_space: u64,

    /// Used disk space in bytes for a specified mount point
    used_disks_space: u64,
}

impl SystemMetrics {
    /// Creates a new `SystemMetrics` instance, refreshing the system information
    async fn new(disk_mount_points: &[PathBuf]) -> Self {
        let mut sys = System::new_all();

        // CPU
        sys.refresh_cpu_usage();
        let mut cpu_usage = sys.global_cpu_usage();
        tokio::time::sleep(sysinfo::MINIMUM_CPU_UPDATE_INTERVAL).await;
        sys.refresh_cpu_usage();
        cpu_usage += sys.global_cpu_usage();
        cpu_usage /= 2.0;

        // Memory
        sys.refresh_memory();
        let total_memory = sys.total_memory();
        let used_memory = sys.used_memory();

        // Swap
        let total_swap = sys.total_swap();
        let used_swap = sys.used_swap();

        // Disks
        let disks = Disks::new_with_refreshed_list();
        let mut total_disks_space = 0;
        let mut used_disks_space = 0;
        for disk in &disks {
            if disk_mount_points.contains(&disk.mount_point().to_path_buf()) {
                total_disks_space += disk.total_space();
                used_disks_space += disk.total_space() - disk.available_space();
            }
        }

        Self {
            cpu_usage,
            total_memory,
            used_memory,
            total_swap,
            used_swap,
            total_disks_space,
            used_disks_space,
        }
    }

    /// Adds the system metrics to Prometheus gauges
    fn add_metrics(&self, service_name: String) {
        gauge!("system_cpu_usage", "service" => service_name.clone()).set(self.cpu_usage);
        gauge!("system_total_memory", "service" => service_name.clone()).set(self.total_memory as f64);
        gauge!("system_used_memory", "service" => service_name.clone()).set(self.used_memory as f64);
        gauge!("system_total_swap", "service" => service_name.clone()).set(self.total_swap as f64);
        gauge!("system_used_swap", "service" => service_name.clone()).set(self.used_swap as f64);
        gauge!("system_total_disks_space", "service" => service_name.clone()).set(self.total_disks_space as f64);
        gauge!("system_used_disks_usage", "service" => service_name).set(self.used_disks_space as f64);
    }
}

impl fmt::Display for SystemMetrics {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "CPUs:       {:.1}%\n\
             Memory:     {} / {}\n\
             Swap:       {} / {}\n\
             Disk usage: {} / {}",
            self.cpu_usage,
            ByteSize::b(self.used_memory),
            ByteSize::b(self.total_memory),
            ByteSize::b(self.used_swap),
            ByteSize::b(self.total_swap),
            ByteSize::b(self.used_disks_space),
            ByteSize::b(self.total_disks_space),
        )
    }
}