Skip to main content

relay_server/services/
autoscaling.rs

1use std::sync::Arc;
2
3use crate::MemoryStat;
4use crate::services::buffer::PartitionedEnvelopeBuffer;
5use crate::services::processor::EnvelopeProcessorServicePool;
6use relay_system::{
7    AsyncResponse, Controller, FromMessage, Handle, Interface, RuntimeMetrics, Sender, Service,
8};
9use tokio::time::Instant;
10
11/// Service that tracks internal relay metrics so that they can be exposed.
12pub struct AutoscalingMetricService {
13    /// For exposing internal memory usage of relay.
14    memory_stat: MemoryStat,
15    /// Reference to the spooler to get item count and total used size.
16    envelope_buffer: Arc<PartitionedEnvelopeBuffer>,
17    /// Runtime handle to expose service utilization metrics.
18    handle: Handle,
19    /// Gives access to runtime metrics.
20    runtime_metrics: RuntimeMetrics,
21    /// The last time the runtime utilization was checked.
22    last_runtime_check: Instant,
23    /// This will always report `1` unless the instance is shutting down.
24    up: u8,
25    /// Gives access to AsyncPool metrics.
26    async_pool: EnvelopeProcessorServicePool,
27}
28
29impl AutoscalingMetricService {
30    pub fn new(
31        memory_stat: MemoryStat,
32        envelope_buffer: Arc<PartitionedEnvelopeBuffer>,
33        handle: Handle,
34        async_pool: EnvelopeProcessorServicePool,
35    ) -> Self {
36        let runtime_metrics = handle.metrics();
37        Self {
38            memory_stat,
39            envelope_buffer,
40            handle,
41            runtime_metrics,
42            last_runtime_check: Instant::now(),
43            async_pool,
44            up: 1,
45        }
46    }
47}
48
49impl Service for AutoscalingMetricService {
50    type Interface = AutoscalingMetrics;
51
52    async fn run(mut self, mut rx: relay_system::Receiver<Self::Interface>) {
53        let mut shutdown = Controller::shutdown_handle();
54        loop {
55            tokio::select! {
56                _ = shutdown.notified() => {
57                    self.up = 0;
58                },
59                Some(message) = rx.recv() => {
60                    match message {
61                        AutoscalingMetrics::Check(sender) => {
62                            let memory_usage = self.memory_stat.memory();
63                            let metrics = self.handle
64                                .current_services_metrics()
65                                .iter()
66                                .map(|(id, metric)| ServiceUtilization {
67                                    name: id.name(),
68                                    instance_id: id.instance_id(),
69                                    utilization: metric.utilization
70                                }
71                            )
72                                .collect();
73                            let worker_pool_utilization = self.async_pool.metrics().total_utilization();
74                            let runtime_utilization = self.runtime_utilization();
75
76                            sender.send(AutoscalingData {
77                                memory_usage: memory_usage.used_percent(),
78                                up: self.up,
79                                total_size: self.envelope_buffer.total_storage_size(),
80                                item_count: self.envelope_buffer.item_count(),
81                                services_metrics: metrics,
82                                worker_pool_utilization,
83                                runtime_utilization
84                            });
85                        }
86                    }
87                }
88            }
89        }
90    }
91}
92
93impl AutoscalingMetricService {
94    fn runtime_utilization(&mut self) -> u8 {
95        let last_checked = self.last_runtime_check.elapsed().as_secs_f64();
96        // Prevent division by 0 in case it's checked in rapid succession.
97        if last_checked < 0.001 {
98            return 0;
99        }
100        let avg_utilization = (0..self.runtime_metrics.num_workers())
101            .map(|worker_id| self.runtime_metrics.worker_total_busy_duration(worker_id))
102            .map(|busy| busy.as_secs_f64())
103            .sum::<f64>()
104            / last_checked
105            / (self.runtime_metrics.num_workers() as f64);
106
107        self.last_runtime_check = Instant::now();
108
109        (avg_utilization * 100.0).min(100.0) as u8
110    }
111}
112
113/// Supported operations within the internal metrics service.
114pub enum AutoscalingMessageKind {
115    /// Requests the current data from the service.
116    Check,
117}
118
119/// This mirrors the same messages as [`AutoscalingMessageKind`] but it can be augmented
120/// with additional data necessary for the service framework, for example a Sender.
121pub enum AutoscalingMetrics {
122    Check(Sender<AutoscalingData>),
123}
124
125impl Interface for AutoscalingMetrics {}
126
127impl FromMessage<AutoscalingMessageKind> for AutoscalingMetrics {
128    type Response = AsyncResponse<AutoscalingData>;
129
130    fn from_message(message: AutoscalingMessageKind, sender: Sender<AutoscalingData>) -> Self {
131        match message {
132            AutoscalingMessageKind::Check => AutoscalingMetrics::Check(sender),
133        }
134    }
135}
136
137/// Contains data that is used for autoscaling.
138pub struct AutoscalingData {
139    /// Memory usage of relay.
140    pub memory_usage: f32,
141    /// Is `1` if relay is running, `0` if it's shutting down.
142    pub up: u8,
143    /// The total number of bytes used by the spooler.
144    pub total_size: u64,
145    /// The total number of envelopes in the spooler.
146    pub item_count: u64,
147    /// Worker pool utilization in percent.
148    pub worker_pool_utilization: u8,
149    /// List of service utilization.
150    pub services_metrics: Vec<ServiceUtilization>,
151    /// Utilization of the async runtime.
152    pub runtime_utilization: u8,
153}
154
155/// Contains the minimal required information for service utilization.
156///
157/// A service can have multiple instances which will all have the same name.
158/// Those instances are distinguished by the `instance_id`.
159pub struct ServiceUtilization {
160    /// The service name.
161    pub name: &'static str,
162    /// The id of the specific service instance.
163    pub instance_id: u32,
164    /// Utilization as percentage.
165    pub utilization: u8,
166}
167
168impl ServiceUtilization {
169    pub fn new(name: &'static str, instance_id: u32, utilization: u8) -> Self {
170        Self {
171            name,
172            instance_id,
173            utilization,
174        }
175    }
176}