rd_util/
systemd.rs

1// Copyright (c) Facebook, Inc. and its affiliates.
2use anyhow::{bail, Result};
3use log::{debug, info, trace, warn};
4use rustbus::{
5    message_builder::{self, MarshalledMessage},
6    params, signature, standard_messages, DuplexConn, MessageType, RpcConn,
7};
8use std::cell::RefCell;
9use std::collections::HashMap;
10use std::fmt;
11use std::sync::atomic::{AtomicU64, Ordering};
12use std::thread::{sleep, LocalKey};
13use std::time::{Duration, Instant};
14
15pub const SYSTEMD_DFL_TIMEOUT: f64 = 15.0;
16const SD1_DST: &str = "org.freedesktop.systemd1";
17const SD1_PATH: &str = "/org/freedesktop/systemd1";
18
19std::thread_local!(pub static SYS_SD_BUS: RefCell<SystemdDbus> =
20                   RefCell::new(SystemdDbus::new(false).unwrap()));
21std::thread_local!(pub static USR_SD_BUS: RefCell<SystemdDbus> =
22                   RefCell::new(SystemdDbus::new(true).unwrap()));
23
24// The following and the explicit error wrappings can be removed once
25// rustbus error implements std::error::Error.
26type RbResult<T> = Result<T, rustbus::connection::Error>;
27
28lazy_static::lazy_static! {
29    static ref SYSTEMD_TIMEOUT_MS: AtomicU64 =
30        AtomicU64::new((SYSTEMD_DFL_TIMEOUT * 1000.0).round() as u64);
31}
32
33pub fn set_systemd_timeout(timeout: f64) {
34    SYSTEMD_TIMEOUT_MS.store((timeout * 1000.0).round() as u64, Ordering::Relaxed);
35}
36
37fn systemd_timeout() -> f64 {
38    SYSTEMD_TIMEOUT_MS.load(Ordering::Relaxed) as f64 / 1000.0
39}
40
41fn systemd_conn_timeout() -> rustbus::connection::Timeout {
42    rustbus::connection::Timeout::Duration(Duration::from_secs_f64(systemd_timeout()))
43}
44
45fn wrap_rustbus_result<T>(r: RbResult<T>) -> Result<T> {
46    match r {
47        Ok(r) => Ok(r),
48        Err(e) => bail!("{:?}", &e),
49    }
50}
51
52#[derive(Debug)]
53pub enum Prop {
54    U32(u32),
55    U64(u64),
56    Bool(bool),
57    String(String),
58}
59
60// define the variant with a fitting marshal and unmarshal impl
61rustbus::dbus_variant_sig!(PropVariant,
62                           Bool => bool;
63                           U32 => u32;
64                           U64 => u64;
65                           String => String;
66                           StringList => Vec<String>;
67                           ExecStart => Vec<(String, Vec<String>, bool)>);
68
69fn dbus_sig(input: &str) -> signature::Type {
70    signature::Type::parse_description(input).as_ref().unwrap()[0].clone()
71}
72
73fn dbus_param_array<'a, 'e>(v: params::Array<'a, 'e>) -> params::Param<'a, 'e> {
74    params::Param::Container(params::Container::Array(v))
75}
76
77fn escape_name(name: &str) -> String {
78    let mut escaped = String::new();
79    for c in name.chars() {
80        let mut buf = [0; 1]; // must be ascii
81        let utf8 = c.encode_utf8(&mut buf);
82
83        if c.is_alphanumeric() {
84            escaped += utf8;
85        } else {
86            escaped += &format!("_{:02x}", utf8.bytes().next().unwrap());
87        }
88    }
89    if log::max_level() >= log::LevelFilter::Trace && name != escaped {
90        trace!("svc: escaped {:?} -> {:?}", &name, &escaped);
91    }
92    escaped
93}
94
95fn systemd_unit_call(method: &str, intf: &str, name: &str) -> MarshalledMessage {
96    let path = SD1_PATH.to_string() + "/unit/" + &escape_name(&name);
97
98    message_builder::MessageBuilder::new()
99        .call(method)
100        .with_interface(intf)
101        .on(path)
102        .at(SD1_DST)
103        .build()
104}
105
106fn systemd_sd1_call(method: &str) -> MarshalledMessage {
107    message_builder::MessageBuilder::new()
108        .call(method)
109        .with_interface("org.freedesktop.systemd1.Manager")
110        .on(SD1_PATH)
111        .at(SD1_DST)
112        .build()
113}
114
115fn systemd_start_transient_svc_call(
116    name: String,
117    args: Vec<String>,
118    envs: Vec<String>,
119    extra_props: Vec<(String, PropVariant)>,
120) -> MarshalledMessage {
121    // NAME(s) JOB_MODE(s) PROPS(a(sv)) AUX_UNITS(a(s a(sv)))
122    //
123    // PROPS:
124    // ["Description"] = str,
125    // ["Slice"] = str,
126    // ["CPUWeight"] = num,
127    // ...
128    // ["Environment"] = ([ENV0]=str, [ENV1]=str...)
129    // ["ExecStart"] = (args[0], (args[0], args[1], ...), false)
130    let mut call = systemd_sd1_call("StartTransientUnit");
131
132    // name and job_mode
133    call.body.push_param2(&name, "fail").unwrap();
134
135    // desc string
136    let desc = args.iter().fold(name.clone(), |mut a, i| {
137        a += " ";
138        a += i;
139        a
140    });
141
142    let mut props = vec![
143        ("Description".to_owned(), PropVariant::String(desc)),
144        ("Environment".to_owned(), PropVariant::StringList(envs)),
145        (
146            "ExecStart".to_owned(),
147            PropVariant::ExecStart(vec![(args[0].clone(), args, false)]),
148        ),
149    ];
150
151    for prop in extra_props.into_iter() {
152        props.push(prop);
153    }
154
155    // assemble props
156    call.body.push_param(props).unwrap();
157
158    // no aux units
159    call.body
160        .push_old_param(&dbus_param_array(params::Array {
161            element_sig: dbus_sig("(sa(sv))"),
162            values: vec![],
163        }))
164        .unwrap();
165
166    call
167}
168
169pub struct SystemdDbus {
170    pub rpc_conn: RpcConn,
171}
172
173impl SystemdDbus {
174    fn new_int(user: bool) -> RbResult<SystemdDbus> {
175        let mut rpc_conn = RpcConn::new(match user {
176            false => DuplexConn::connect_to_bus(rustbus::get_system_bus_path()?, true)?,
177            true => DuplexConn::connect_to_bus(rustbus::get_session_bus_path()?, true)?,
178        });
179
180        rpc_conn.set_filter(Box::new(|msg| match msg.typ {
181            MessageType::Error => true,
182            MessageType::Reply => true,
183            _ => false,
184        }));
185
186        let mut sysdbus = Self { rpc_conn };
187        sysdbus.send_msg_and_wait_int(&mut standard_messages::hello())?;
188        Ok(sysdbus)
189    }
190
191    pub fn new(user: bool) -> Result<SystemdDbus> {
192        wrap_rustbus_result(Self::new_int(user))
193    }
194
195    fn send_msg_and_wait_int(
196        &mut self,
197        msg: &mut MarshalledMessage,
198    ) -> RbResult<MarshalledMessage> {
199        let msg_serial = self.rpc_conn.send_message(msg)?.write_all().unwrap();
200        self.rpc_conn
201            .wait_response(msg_serial, systemd_conn_timeout())
202    }
203
204    pub fn send_msg_and_wait(&mut self, msg: &mut MarshalledMessage) -> Result<MarshalledMessage> {
205        wrap_rustbus_result(self.send_msg_and_wait_int(msg))
206    }
207
208    pub fn daemon_reload(&mut self) -> Result<()> {
209        let mut msg = systemd_sd1_call("Reload");
210        self.send_msg_and_wait(&mut msg)?;
211        Ok(())
212    }
213
214    pub fn get_unit_props<'u>(&mut self, name: &str) -> Result<params::Param<'static, 'static>> {
215        let mut msg = systemd_unit_call("GetAll", "org.freedesktop.DBus.Properties", name);
216        msg.body.push_param("").unwrap();
217        let resp = match self.send_msg_and_wait(&mut msg)?.unmarshall_all() {
218            Ok(v) => v,
219            Err(e) => bail!("failed to unmarshall GetAll response ({:?})", &e),
220        };
221        match resp.params.into_iter().next() {
222            Some(props) => Ok(props),
223            None => bail!("GetAll response doesn't have any data"),
224        }
225    }
226
227    pub fn set_unit_props(&mut self, name: &str, props: Vec<(String, PropVariant)>) -> Result<()> {
228        let mut msg = systemd_sd1_call("SetUnitProperties");
229        msg.body.push_param3(name, true, props).unwrap();
230        self.send_msg_and_wait(&mut msg)?;
231        Ok(())
232    }
233
234    pub fn start_unit(&mut self, name: &str) -> Result<()> {
235        let mut msg = systemd_sd1_call("StartUnit");
236        msg.body.push_param2(&name, "fail").unwrap();
237        self.send_msg_and_wait(&mut msg)?;
238        Ok(())
239    }
240
241    pub fn stop_unit(&mut self, name: &str) -> Result<()> {
242        let mut msg = systemd_sd1_call("StopUnit");
243        msg.body.push_param2(&name, "fail").unwrap();
244        self.send_msg_and_wait(&mut msg)?;
245        Ok(())
246    }
247
248    pub fn reset_failed_unit(&mut self, name: &str) -> Result<()> {
249        let mut msg = systemd_sd1_call("ResetFailedUnit");
250        msg.body.push_param(&name).unwrap();
251        self.send_msg_and_wait(&mut msg)?;
252        Ok(())
253    }
254
255    pub fn restart_unit(&mut self, name: &str) -> Result<()> {
256        let mut msg = systemd_sd1_call("RestartUnit");
257        msg.body.push_param2(&name, "fail").unwrap();
258        self.send_msg_and_wait(&mut msg)?;
259        Ok(())
260    }
261
262    pub fn start_transient_svc(
263        &mut self,
264        name: String,
265        args: Vec<String>,
266        envs: Vec<String>,
267        extra_props: Vec<(String, PropVariant)>,
268    ) -> Result<()> {
269        let mut msg = systemd_start_transient_svc_call(name, args, envs, extra_props);
270        self.send_msg_and_wait(&mut msg)?;
271        Ok(())
272    }
273}
274
275pub fn daemon_reload() -> Result<()> {
276    SYS_SD_BUS.with(|s| s.borrow_mut().daemon_reload())
277}
278
279#[derive(Debug, Clone, PartialEq, Eq)]
280pub enum UnitState {
281    NotFound,
282    Running,
283    Exited,
284    OtherActive(String),
285    Inactive(String),
286    Failed(String),
287    Other(String),
288}
289
290use UnitState as US;
291
292impl Default for UnitState {
293    fn default() -> Self {
294        US::NotFound
295    }
296}
297
298#[derive(Debug)]
299pub struct UnitProps {
300    props: HashMap<String, Prop>,
301}
302
303impl UnitProps {
304    fn new(dict: &params::Param) -> Result<Self> {
305        let dict = match dict {
306            params::Param::Container(params::Container::Dict(dict)) => dict,
307            _ => bail!("dict type mismatch"),
308        };
309
310        let mut props = HashMap::new();
311
312        for (k, v) in dict.map.iter() {
313            if let (
314                params::Base::String(key),
315                params::Param::Container(params::Container::Variant(boxed)),
316            ) = (k, v)
317            {
318                match &boxed.value {
319                    params::Param::Base(params::Base::String(v)) => {
320                        props.insert(key.into(), Prop::String(v.into()));
321                    }
322                    params::Param::Base(params::Base::Uint32(v)) => {
323                        props.insert(key.into(), Prop::U32(*v));
324                    }
325                    params::Param::Base(params::Base::Uint64(v)) => {
326                        props.insert(key.into(), Prop::U64(*v));
327                    }
328                    params::Param::Base(params::Base::Boolean(v)) => {
329                        props.insert(key.into(), Prop::Bool(*v));
330                    }
331                    _ => {}
332                }
333            }
334        }
335        Ok(Self { props })
336    }
337
338    pub fn string<'a>(&'a self, key: &str) -> Option<&'a str> {
339        match self.props.get(key) {
340            Some(Prop::String(v)) => Some(v),
341            _ => None,
342        }
343    }
344
345    pub fn u64_dfl_max(&self, key: &str) -> Option<u64> {
346        match self.props.get(key) {
347            Some(Prop::U64(v)) if *v < std::u64::MAX => Some(*v),
348            _ => None,
349        }
350    }
351
352    pub fn u64_dfl_zero(&self, key: &str) -> Option<u64> {
353        match self.props.get(key) {
354            Some(Prop::U64(v)) if *v > 0 => Some(*v),
355            _ => None,
356        }
357    }
358
359    fn state(&self) -> US {
360        let v = self.string("LoadState");
361        match v {
362            Some("loaded") => {}
363            Some("not-found") => return US::NotFound,
364            Some(_) => return US::Other(v.unwrap().into()),
365            None => return US::Other("no-load-state".into()),
366        };
367
368        let ss = match self.string("SubState") {
369            Some(v) => v.to_string(),
370            None => "no-sub-state".to_string(),
371        };
372
373        match self.string("ActiveState") {
374            Some("active") => match ss.as_str() {
375                "running" => US::Running,
376                "exited" => US::Exited,
377                _ => US::OtherActive(ss),
378            },
379            Some("inactive") => US::Inactive(ss.into()),
380            Some("failed") => US::Failed(ss.into()),
381            Some(v) => US::Other(format!("{}:{}", v, ss)),
382            None => US::Other("no-active-state".into()),
383        }
384    }
385}
386
387#[derive(Debug, Default, Clone, PartialEq, Eq)]
388pub struct UnitResCtl {
389    pub cpu_weight: Option<u64>,
390    pub io_weight: Option<u64>,
391    pub mem_min: Option<u64>,
392    pub mem_low: Option<u64>,
393    pub mem_high: Option<u64>,
394    pub mem_max: Option<u64>,
395}
396
397impl fmt::Display for UnitResCtl {
398    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
399        write!(
400            f,
401            "cpu_w={:?} io_w={:?} mem={:?}:{:?}:{:?}:{:?}",
402            &self.cpu_weight,
403            &self.io_weight,
404            &self.mem_min,
405            &self.mem_low,
406            &self.mem_high,
407            &self.mem_max
408        )
409    }
410}
411
412#[derive(Debug)]
413pub struct Unit {
414    pub user: bool,
415    pub name: String,
416    pub state: US,
417    pub resctl: UnitResCtl,
418    pub props: UnitProps,
419    pub quiet: bool,
420}
421
422impl fmt::Display for Unit {
423    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
424        let user_str = match self.user {
425            true => "(user)",
426            false => "",
427        };
428        write!(
429            f,
430            "{}{}: state={:?} {}",
431            &self.name, &user_str, &self.state, &self.resctl,
432        )
433    }
434}
435
436impl Unit {
437    pub fn new(user: bool, name: String) -> Result<Self> {
438        let sb = match user {
439            false => &SYS_SD_BUS,
440            true => &USR_SD_BUS,
441        };
442        let mut svc = Self {
443            user,
444            state: US::Other("unknown".into()),
445            resctl: Default::default(),
446            props: UnitProps::new(&(sb.with(|s| s.borrow_mut().get_unit_props(&name))?))?,
447            quiet: false,
448            name,
449        };
450        svc.refresh_fields();
451        Ok(svc)
452    }
453
454    pub fn new_sys(name: String) -> Result<Self> {
455        Self::new(false, name)
456    }
457
458    pub fn new_user(name: String) -> Result<Self> {
459        Self::new(true, name)
460    }
461
462    pub fn sd_bus(&self) -> &'static LocalKey<RefCell<SystemdDbus>> {
463        match self.user {
464            false => &SYS_SD_BUS,
465            true => &USR_SD_BUS,
466        }
467    }
468
469    pub fn refresh(&mut self) -> Result<()> {
470        trace!("svc: {:?} refreshing", &self.name);
471        self.props = match self
472            .sd_bus()
473            .with(|s| s.borrow_mut().get_unit_props(&self.name))
474        {
475            Ok(props) => UnitProps::new(&props)?,
476            Err(e) => {
477                debug!(
478                    "Failed to unmarshall response from {}, assuming gone ({:?})",
479                    &self.name, &e
480                );
481                self.state = US::NotFound;
482                return Err(e);
483            }
484        };
485        self.refresh_fields();
486        Ok(())
487    }
488
489    pub fn refresh_fields(&mut self) {
490        let new_state = self.props.state();
491
492        if !self.quiet && self.state == US::Running {
493            match &new_state {
494                US::Running => {}
495                US::Exited => info!("svc: {:?} exited", &self.name),
496                US::Failed(how) => info!("svc: {:?} failed ({:?})", &self.name, &how),
497                US::NotFound => info!("svc: {:?} is gone", &self.name),
498                s => info!(
499                    "svc: {:?} transitioned from Running to {:?}",
500                    &self.name, &s
501                ),
502            }
503        }
504
505        self.state = new_state;
506        self.resctl.cpu_weight = self.props.u64_dfl_max("CPUWeight");
507        self.resctl.io_weight = self.props.u64_dfl_max("IOWeight");
508        self.resctl.mem_min = self.props.u64_dfl_zero("MemoryMin");
509        self.resctl.mem_low = self.props.u64_dfl_zero("MemoryLow");
510        self.resctl.mem_high = self.props.u64_dfl_max("MemoryHigh");
511        self.resctl.mem_max = self.props.u64_dfl_max("MemoryMax");
512    }
513
514    pub fn resctl_props(&self) -> Vec<(String, PropVariant)> {
515        vec![
516            (
517                "CPUWeight".into(),
518                PropVariant::U64(self.resctl.cpu_weight.unwrap_or(u64::MAX)),
519            ),
520            (
521                "IOWeight".into(),
522                PropVariant::U64(self.resctl.io_weight.unwrap_or(u64::MAX)),
523            ),
524            (
525                "MemoryMin".into(),
526                PropVariant::U64(self.resctl.mem_min.unwrap_or(0)),
527            ),
528            (
529                "MemoryLow".into(),
530                PropVariant::U64(self.resctl.mem_low.unwrap_or(0)),
531            ),
532            (
533                "MemoryHigh".into(),
534                PropVariant::U64(self.resctl.mem_high.unwrap_or(std::u64::MAX)),
535            ),
536            (
537                "MemoryMax".into(),
538                PropVariant::U64(self.resctl.mem_max.unwrap_or(std::u64::MAX)),
539            ),
540        ]
541    }
542
543    pub fn apply(&mut self) -> Result<()> {
544        trace!("svc: {:?} applying resctl", &self.name);
545        self.sd_bus().with(|s| {
546            s.borrow_mut()
547                .set_unit_props(&self.name, self.resctl_props())
548        })?;
549        self.refresh()
550    }
551
552    pub fn set_prop(&mut self, key: &str, prop: Prop) -> Result<()> {
553        let props = match prop {
554            Prop::U32(v) => PropVariant::U32(v),
555            Prop::U64(v) => PropVariant::U64(v),
556            Prop::Bool(v) => PropVariant::Bool(v),
557            Prop::String(v) => PropVariant::String(v),
558        };
559        self.sd_bus().with(|s| {
560            s.borrow_mut()
561                .set_unit_props(&self.name, vec![(key.into(), props)])
562        })?;
563        self.refresh()
564    }
565
566    fn wait_transition<F>(&mut self, wait_till: F, timeout: f64, exiting_timeout: f64)
567    where
568        F: Fn(&US) -> bool,
569    {
570        let started_at = Instant::now();
571        loop {
572            if let Ok(()) = self.refresh() {
573                trace!(
574                    "svc: {:?} waiting transitions ({:?})",
575                    &self.name,
576                    &self.state
577                );
578                match &self.state {
579                    US::OtherActive(_) | US::Other(_) => {}
580                    state if !wait_till(state) => {}
581                    _ => return,
582                }
583            }
584
585            let dur = Duration::from_secs_f64(match super::prog_exiting() {
586                false => timeout,
587                true => exiting_timeout,
588            });
589            if Instant::now().duration_since(started_at) >= dur {
590                trace!("svc: {:?} waiting transitions timed out", &self.name);
591                return;
592            }
593
594            sleep(Duration::from_millis(100));
595        }
596    }
597
598    pub fn stop(&mut self) -> Result<bool> {
599        debug!("svc: {:?} stopping ({:?})", &self.name, &self.state);
600
601        self.refresh()?;
602        match self.state {
603            US::NotFound | US::Failed(_) => {
604                debug!("svc: {:?} already stopped ({:?})", &self.name, &self.state);
605                return Ok(true);
606            }
607            _ => {}
608        }
609
610        self.sd_bus()
611            .with(|s| s.borrow_mut().stop_unit(&self.name))?;
612        // We're used from exit paths. Force a bit of wait so that we
613        // can shut down gracefully in most cases.
614        self.wait_transition(
615            |x| *x != US::Running,
616            systemd_timeout(),
617            systemd_timeout() / 5.0,
618        );
619        if !self.quiet {
620            info!("svc: {:?} stopped ({:?})", &self.name, &self.state);
621        }
622        match self.state {
623            US::NotFound | US::Failed(_) => Ok(true),
624            _ => Ok(false),
625        }
626    }
627
628    pub fn stop_and_reset(&mut self) -> Result<()> {
629        self.stop()?;
630        if let US::Failed(_) = self.state {
631            self.sd_bus()
632                .with(|s| s.borrow_mut().reset_failed_unit(&self.name))?;
633            // We're used from exit paths. Force a bit of wait so that we
634            // can shut down gracefully in most cases.
635            self.wait_transition(
636                |x| *x == US::NotFound,
637                systemd_timeout(),
638                systemd_timeout() / 5.0,
639            );
640        }
641        match self.state {
642            US::NotFound => Ok(()),
643            _ => bail!(
644                "invalid post-reset state {:?} for {}",
645                self.state,
646                &self.name
647            ),
648        }
649    }
650
651    pub fn try_start_nowait(&mut self) -> Result<()> {
652        debug!("svc: {:?} starting ({:?})", &self.name, &self.state);
653        self.sd_bus()
654            .with(|s| s.borrow_mut().start_unit(&self.name))
655    }
656
657    pub fn try_start(&mut self) -> Result<bool> {
658        self.try_start_nowait()?;
659        self.wait_transition(
660            |x| match x {
661                US::Running | US::Exited | US::Failed(_) => true,
662                _ => false,
663            },
664            systemd_timeout(),
665            0.0,
666        );
667        if !self.quiet {
668            info!("svc: {:?} started ({:?})", &self.name, &self.state);
669        }
670        match self.state {
671            US::Running | US::Exited => Ok(true),
672            _ => Ok(false),
673        }
674    }
675
676    pub fn restart(&mut self) -> Result<()> {
677        if !self.quiet {
678            info!("svc: {:?} restarting ({:?})", &self.name, &self.state);
679        }
680        self.sd_bus()
681            .with(|s| s.borrow_mut().restart_unit(&self.name))
682    }
683}
684
685pub struct TransientService {
686    pub unit: Unit,
687    pub args: Vec<String>,
688    pub envs: Vec<String>,
689    pub extra_props: HashMap<String, Prop>,
690    pub keep: bool,
691}
692
693impl TransientService {
694    pub fn new(
695        user: bool,
696        name: String,
697        args: Vec<String>,
698        envs: Vec<String>,
699        umask: Option<u32>,
700    ) -> Result<Self> {
701        if !name.ends_with(".service") {
702            bail!("invalid service name {}", &name);
703        }
704        let mut svc = Self {
705            unit: Unit::new(user, name)?,
706            args,
707            envs,
708            extra_props: HashMap::new(),
709            keep: false,
710        };
711        svc.add_prop("RemainAfterExit".into(), Prop::Bool(true));
712        if let Some(v) = umask {
713            svc.add_prop("UMask".into(), Prop::U32(v));
714        }
715        Ok(svc)
716    }
717
718    pub fn new_sys(
719        name: String,
720        args: Vec<String>,
721        envs: Vec<String>,
722        umask: Option<u32>,
723    ) -> Result<Self> {
724        Self::new(false, name, args, envs, umask)
725    }
726
727    pub fn new_user(
728        name: String,
729        args: Vec<String>,
730        envs: Vec<String>,
731        umask: Option<u32>,
732    ) -> Result<Self> {
733        Self::new(true, name, args, envs, umask)
734    }
735
736    pub fn add_prop(&mut self, key: String, v: Prop) -> &mut Self {
737        self.extra_props.insert(key, v);
738        self
739    }
740
741    pub fn del_prop(&mut self, key: &String) -> (&mut Self, Option<Prop>) {
742        let v = self.extra_props.remove(key);
743        (self, v)
744    }
745
746    pub fn set_slice(&mut self, slice: &str) -> &mut Self {
747        self.add_prop("Slice".into(), Prop::String(slice.into()));
748        self
749    }
750
751    pub fn set_working_dir(&mut self, dir: &str) -> &mut Self {
752        self.add_prop("WorkingDirectory".into(), Prop::String(dir.into()));
753        self
754    }
755
756    pub fn set_restart_always(&mut self) -> &mut Self {
757        self.add_prop("Restart".into(), Prop::String("always".into()));
758        self
759    }
760
761    pub fn set_quiet(&mut self) -> &mut Self {
762        self.unit.quiet = true;
763        self
764    }
765
766    fn try_start(&mut self) -> Result<bool> {
767        let mut extra_props = self.unit.resctl_props();
768        for (k, v) in self.extra_props.iter() {
769            let variant = match v {
770                Prop::U32(v) => PropVariant::U32(*v),
771                Prop::U64(v) => PropVariant::U64(*v),
772                Prop::Bool(v) => PropVariant::Bool(*v),
773                Prop::String(v) => PropVariant::String(v.clone()),
774            };
775            extra_props.push((k.clone(), variant));
776        }
777
778        debug!(
779            "svc: {:?} starting ({:?})",
780            &self.unit.name, &self.unit.state
781        );
782        self.unit.sd_bus().with(|s| {
783            s.borrow_mut().start_transient_svc(
784                self.unit.name.clone(),
785                self.args.clone(),
786                self.envs.clone(),
787                extra_props,
788            )
789        })?;
790
791        self.unit.wait_transition(
792            |x| match x {
793                US::Running | US::Exited | US::Failed(_) => true,
794                _ => false,
795            },
796            systemd_timeout(),
797            0.0,
798        );
799        if !self.unit.quiet {
800            info!(
801                "svc: {:?} started ({:?})",
802                &self.unit.name, &self.unit.state
803            );
804        }
805        match self.unit.state {
806            US::Running | US::Exited => Ok(true),
807            _ => Ok(false),
808        }
809    }
810
811    pub fn start(&mut self) -> Result<()> {
812        let resctl = self.unit.resctl.clone();
813        match self.unit.stop_and_reset() {
814            Ok(()) => {
815                self.unit.resctl = resctl;
816                match self.try_start() {
817                    Ok(true) => Ok(()),
818                    Ok(false) => bail!("invalid service state {:?}", &self.unit.state),
819                    Err(e) => Err(e),
820                }
821            }
822            Err(e) => Err(e),
823        }
824    }
825}
826
827impl Drop for TransientService {
828    fn drop(&mut self) {
829        if self.keep {
830            return;
831        }
832        for tries in (1..6).rev() {
833            let action = match tries {
834                0 => String::new(),
835                v => format!(", retrying... ({} tries left)", v),
836            };
837            match self.unit.stop_and_reset() {
838                Ok(()) => {}
839                Err(e) => warn!(
840                    "Failed to stop {} on drop ({:?}){}",
841                    &self.unit.name, &e, action
842                ),
843            }
844        }
845    }
846}
847
848#[cfg(test)]
849mod tests {
850    use super::{TransientService, UnitState};
851    use log::{info, trace};
852    use std::thread::sleep;
853    use std::time::Duration;
854
855    //#[test]
856    // TODO: This test is not hermetic as it has an implicit dependency
857    // on the systemd session bus; it should be spinning up its own bus instead.
858    #[allow(dead_code)]
859    fn test_transient_service() {
860        let _ = ::env_logger::try_init();
861        let name = "test-transient.service";
862
863        info!("Creating {}", &name);
864        let mut svc = TransientService::new_user(
865            name.into(),
866            vec![
867                "/usr/bin/bash".into(),
868                "-c".into(),
869                "echo $TEST_ENV; sleep 3".into(),
870            ],
871            vec![("TEST_ENV=TEST".into())],
872            None,
873        )
874        .unwrap();
875        assert_eq!(svc.unit.state, UnitState::NotFound);
876
877        info!("Starting the service");
878        svc.start().unwrap();
879
880        trace!("{} props={:#?}", &name, &svc.unit.props);
881        info!("{}", &svc.unit);
882
883        info!("Setting cpu weight to 111");
884        let cpu_weight = svc.unit.resctl.cpu_weight;
885        svc.unit.resctl.cpu_weight = Some(111);
886        svc.unit.apply().unwrap();
887        info!("{}", &svc.unit);
888        assert_eq!(svc.unit.resctl.cpu_weight, Some(111));
889
890        info!("Restoring cpu weight");
891        svc.unit.resctl.cpu_weight = cpu_weight;
892        svc.unit.apply().unwrap();
893        info!("{}", &svc.unit);
894        assert_eq!(svc.unit.resctl.cpu_weight, cpu_weight);
895
896        info!("Sleeping 4 secs and checking state, it should have exited");
897        sleep(Duration::from_secs(4));
898        svc.unit.refresh().unwrap();
899        info!("{}", &svc.unit);
900        assert_eq!(svc.unit.state, UnitState::Exited);
901
902        info!("Restarting the service w/o RemainAfterExit");
903        svc.del_prop(&"RemainAfterExit".to_string());
904        svc.start().unwrap();
905
906        info!("Sleeping 4 secs and checking state, it should be gone");
907        sleep(Duration::from_secs(4));
908        svc.unit.refresh().unwrap();
909        info!("{}", &svc.unit);
910        assert_eq!(svc.unit.state, UnitState::NotFound);
911
912        info!("Dropping the service");
913        drop(svc);
914        info!("Dropped");
915    }
916}
rd_util/systemd.rs

rd_util/
systemd.rs
