relay_server/services/
autoscaling.rs1use 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
11pub struct AutoscalingMetricService {
13 memory_stat: MemoryStat,
15 envelope_buffer: Arc<PartitionedEnvelopeBuffer>,
17 handle: Handle,
19 runtime_metrics: RuntimeMetrics,
21 last_runtime_check: Instant,
23 up: u8,
25 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 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
113pub enum AutoscalingMessageKind {
115 Check,
117}
118
119pub 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
137pub struct AutoscalingData {
139 pub memory_usage: f32,
141 pub up: u8,
143 pub total_size: u64,
145 pub item_count: u64,
147 pub worker_pool_utilization: u8,
149 pub services_metrics: Vec<ServiceUtilization>,
151 pub runtime_utilization: u8,
153}
154
155pub struct ServiceUtilization {
160 pub name: &'static str,
162 pub instance_id: u32,
164 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}