@@ -21,6 +21,7 @@ use prometheus_client::{encoding::text::encode, registry::Registry};
2121use serde:: Deserialize ;
2222use std:: {
2323 borrow:: Cow ,
24+ collections:: BTreeMap ,
2425 io:: { self , BufRead as _, BufReader } ,
2526} ;
2627use tokio:: process:: Command ;
@@ -30,11 +31,31 @@ use tower::{
3031} ;
3132use tower_http:: compression:: CompressionLayer ;
3233
33- #[ derive( Debug , Deserialize ) ]
34- pub struct Params {
35- // Only enable jobstats if "jobstats=true"
36- #[ serde( default ) ]
37- jobstats : bool ,
34+ #[ derive( Debug , Deserialize , PartialEq , PartialOrd , Eq , Ord , Hash , Copy , Clone ) ]
35+ #[ serde( rename_all = "snake_case" ) ]
36+ pub enum Dimension {
37+ Jobstats ,
38+ Lnet ,
39+ Lustre ,
40+ LnetStats ,
41+ }
42+
43+ pub type Params = BTreeMap < Dimension , bool > ;
44+
45+ static DEFAULT_PARAMS : [ ( Dimension , bool ) ; 3 ] = [
46+ ( Dimension :: Lnet , true ) ,
47+ ( Dimension :: Lustre , true ) ,
48+ ( Dimension :: LnetStats , true ) ,
49+ ] ;
50+
51+ trait EnableConvenienceExt {
52+ fn enabled ( & self , param : & Dimension ) -> bool ;
53+ }
54+
55+ impl EnableConvenienceExt for Params {
56+ fn enabled ( & self , param : & Dimension ) -> bool {
57+ self . get ( param) . copied ( ) . unwrap_or_default ( )
58+ }
3859}
3960
4061const TIMEOUT_DURATION_SECS : u64 = 120 ;
@@ -149,11 +170,18 @@ pub fn lnet_stats_output() -> Command {
149170pub async fn scrape ( Query ( params) : Query < Params > ) -> Result < Response < Body > , Error > {
150171 let mut registry = Registry :: default ( ) ;
151172
152- if params. jobstats {
153- let child = tokio:: task:: spawn_blocking ( move || jobstats_metrics_cmd ( ) . spawn ( ) ) . await ?;
173+ let mut targets = BTreeMap :: from ( DEFAULT_PARAMS ) ;
174+ for ( param, value) in params {
175+ targets
176+ . entry ( param)
177+ . and_modify ( |v| * v = value)
178+ . or_insert ( value) ;
179+ }
154180
155- if let Ok ( mut child) =
156- child. inspect_err ( |e| tracing:: debug!( "Error while spawning lctl jobstats: {e}" ) )
181+ if targets. enabled ( & Dimension :: Jobstats ) {
182+ if let Ok ( mut child) = tokio:: task:: spawn_blocking ( move || jobstats_metrics_cmd ( ) . spawn ( ) )
183+ . await ?
184+ . inspect_err ( |e| tracing:: debug!( "Error while spawning lctl jobstats: {e}" ) )
157185 {
158186 let reader = BufReader :: with_capacity (
159187 128 * 1_024 ,
@@ -184,23 +212,28 @@ pub async fn scrape(Query(params): Query<Params>) -> Result<Response<Body>, Erro
184212 . await ?
185213 . register_metric ( & mut registry) ;
186214 }
187- } else {
215+ }
216+
217+ if targets. enabled ( & Dimension :: Lustre ) {
188218 let mut output = vec ! [ ] ;
189219
190220 let lctl = lustre_metrics_output ( ) . output ( ) . await ?;
191-
192221 let mut lctl_output = parse_lctl_output ( & lctl. stdout ) ?;
193222
194223 output. append ( & mut lctl_output) ;
195224
196- let lnetctl = net_show_output ( ) . output ( ) . await ?;
225+ // Build the lustre stats
226+ let mut opentelemetry_metrics = Metrics :: default ( ) ;
197227
198- let mut lnetctl_output = parse_lnetctl_output ( & lnetctl. stdout ) ?;
228+ // Build and register Lustre metrics
229+ metrics:: build_lustre_stats ( & output, & mut opentelemetry_metrics) ;
230+ opentelemetry_metrics. register_metric ( & mut registry) ;
231+ }
199232
200- output. append ( & mut lnetctl_output) ;
233+ if targets. enabled ( & Dimension :: LnetStats ) {
234+ let mut output = vec ! [ ] ;
201235
202236 let lnetctl_stats_output = lnet_stats_output ( ) . output ( ) . await ?;
203-
204237 let mut lnetctl_stats_record = parse_lnetctl_stats ( & lnetctl_stats_output. stdout ) ?;
205238
206239 output. append ( & mut lnetctl_stats_record) ;
@@ -213,6 +246,22 @@ pub async fn scrape(Query(params): Query<Params>) -> Result<Response<Body>, Erro
213246 opentelemetry_metrics. register_metric ( & mut registry) ;
214247 }
215248
249+ if targets. enabled ( & Dimension :: Lnet ) {
250+ let mut output = vec ! [ ] ;
251+
252+ let lnetctl = net_show_output ( ) . output ( ) . await ?;
253+ let mut lnetctl_output = parse_lnetctl_output ( & lnetctl. stdout ) ?;
254+
255+ output. append ( & mut lnetctl_output) ;
256+
257+ // Build the lustre stats
258+ let mut opentelemetry_metrics = Metrics :: default ( ) ;
259+
260+ // Build and register Lustre metrics
261+ metrics:: build_lustre_stats ( & output, & mut opentelemetry_metrics) ;
262+ opentelemetry_metrics. register_metric ( & mut registry) ;
263+ }
264+
216265 let mut buffer = String :: new ( ) ;
217266 encode ( & mut buffer, & registry) ?;
218267
0 commit comments