systemd-transient 0.2.0

use std::time::Duration;

use futures_util::StreamExt;
use tracing::warn;
use zbus::{Connection, ObjectPath, OwnedObjectPath, Variant, proxy::CacheProperties};

use crate::{
  JobMode, ReusableTransientUnit, TransientUnitHandle, UnitExit,
  dbus::{
    manager::{JobRemovedStream, ManagerProxy},
    unit::UnitProxy,
    unit_object_path,
  },
  error::SystemdTransientError,
  name::unit_name,
  props::UnitProperty,
  watch::{ExitWatch, exit_of},
};

const UNIT_UNLOAD_TIMEOUT: Duration = Duration::from_secs(5);
const ERR_NO_SUCH_UNIT: &str = "org.freedesktop.systemd1.NoSuchUnit";

#[derive(Clone, Debug)]
pub struct SystemdTransientManager {
  conn: Connection,
  /// Shared by all calls. Property caching is off: this crate reads no
  /// Manager properties, and a cache costs a GetAll plus a match rule.
  proxy: ManagerProxy<'static>,
}

impl SystemdTransientManager {
  pub async fn system() -> Result<Self, SystemdTransientError> {
    Self::from_connection(Connection::system().await?).await
  }

  pub async fn session() -> Result<Self, SystemdTransientError> {
    Self::from_connection(Connection::session().await?).await
  }

  async fn from_connection(conn: Connection) -> Result<Self, SystemdTransientError> {
    let proxy = ManagerProxy::builder(&conn)
      .cache_properties(CacheProperties::No)
      .build()
      .await?;
    proxy.subscribe().await?;

    Ok(Self { conn, proxy })
  }

  /// See [`start`](Self::start) for `name` and `props`.
  pub fn new_reusable(
    &self,
    name: impl Into<String>,
    props: impl IntoIterator<Item = UnitProperty>,
  ) -> ReusableTransientUnit {
    ReusableTransientUnit {
      name: name.into(),
      props: props.into_iter().collect(),
      unit_handle: None,
      systemd_manager: self.clone(),
    }
  }

  /// Starts `name` as a transient unit and returns once its start job
  /// finished.
  ///
  /// `name` keeps a unit type suffix (`.slice`, `.socket`, `.scope`, ...);
  /// any other name gets `.service` appended. A slice that other units name
  /// in `Slice=` before it is started is created implicitly and is not
  /// transient, so start a transient slice before the units in it.
  pub async fn start(
    &self,
    name: &str,
    props: impl IntoIterator<Item = UnitProperty>,
  ) -> Result<TransientUnitHandle, SystemdTransientError> {
    let name = unit_name(name);

    self.prepare_name(&name).await?;

    // Pin the unit to this connection, so it stays loaded after it exits.
    // Appended last, so it wins over a caller's AddRef.
    let pairs: Vec<_> = props
      .into_iter()
      .chain([UnitProperty::AddRef(true)])
      .map(UnitProperty::into_pair)
      .collect();
    let tuples: Vec<(&str, &Variant<'_>)> = pairs.iter().map(|(k, v)| (k.as_ref(), v)).collect();
    let refs: Vec<&(&str, &Variant<'_>)> = tuples.iter().collect();

    let unit = UnitProxy::builder(&self.conn)
      .path(unit_object_path(&name)?)
      .cache_properties(CacheProperties::No)
      .build()
      .await?;

    // Once this returns Ok, the unit exists and holds our reference, which a
    // later failure must drop.
    let jobs = self.proxy.receive_job_removed().await?;
    let job = self
      .proxy
      .start_transient_unit(&name, JobMode::Replace.as_str(), &refs, &[])
      .await?;

    let started = async {
      wait_for_job(jobs, job.as_ref()).await?;
      let watch = ExitWatch::subscribe(&self.conn, &unit).await?;
      // Identifies this run; a later start of the same unit gets a new one.
      let invocation_id = unit.invocation_id().await?;
      Ok::<_, SystemdTransientError>((watch, invocation_id))
    }
    .await;

    let (watch, invocation_id) = match started {
      Ok(started) => started,
      Err(err) => {
        // Let systemd collect the unit this call created.
        if let Err(unref_err) = unit.unref().await {
          warn!("releasing unit {} failed: {}", name, unref_err);
        }
        return Err(err);
      }
    };

    Ok(TransientUnitHandle {
      name,
      manager: self.clone(),
      unit,
      invocation_id,
      watch,
      released: false,
    })
  }

  /// Takes over the transient unit `name`, for example one that an earlier
  /// process started and did not stop. The handle holds its own reference
  /// and works like one from [`start`](Self::start); the unit keeps the
  /// properties it was started with. `name` is completed as for `start`.
  ///
  /// Running and exited units both qualify, as long as systemd still has
  /// them loaded: a failed unit stays loaded until it is reset, an inactive
  /// one only while something references it. Returns `None` if no unit of
  /// that name is loaded, and `UnitNotTransient` for a unit from a unit file.
  pub async fn adopt(
    &self,
    name: &str,
  ) -> Result<Option<TransientUnitHandle>, SystemdTransientError> {
    let name = unit_name(name);

    let Some(path) = get_unit(&self.proxy, &name).await? else {
      return Ok(None);
    };
    let unit = UnitProxy::builder(&self.conn)
      .path(path)
      .cache_properties(CacheProperties::No)
      .build()
      .await?;

    // Reference first, so the unit cannot be collected while we look at it.
    // If it was collected right after GetUnit, this call loads a stub of the
    // name instead, which the LoadState check catches.
    unit.ref_().await?;

    let adopted = async {
      if unit.load_state().await? != "loaded" {
        return Ok(None);
      }
      if !unit.transient().await? {
        return Err(SystemdTransientError::UnitNotTransient { name: name.clone() });
      }
      let watch = ExitWatch::subscribe(&self.conn, &unit).await?;
      let invocation_id = unit.invocation_id().await?;
      Ok(Some((watch, invocation_id)))
    }
    .await;

    match adopted {
      Ok(Some((watch, invocation_id))) => Ok(Some(TransientUnitHandle {
        name,
        manager: self.clone(),
        unit,
        invocation_id,
        watch,
        released: false,
      })),
      other => {
        if let Err(err) = unit.unref().await {
          warn!("releasing unit {} failed: {}", name, err);
        }
        other.map(|_| None)
      }
    }
  }

  /// Loaded units whose names match one of `patterns`, shell globs as for
  /// `systemctl list-units`. No patterns match every loaded unit.
  pub async fn list_units(
    &self,
    patterns: &[&str],
  ) -> Result<Vec<ListedUnit>, SystemdTransientError> {
    let units = self.proxy.list_units_by_patterns(&[], patterns).await?;
    Ok(
      units
        .into_iter()
        .map(
          |(name, _description, load_state, active_state, sub_state, ..)| ListedUnit {
            name,
            load_state,
            active_state,
            sub_state,
          },
        )
        .collect(),
    )
  }

  /// Makes `name` free for `StartTransientUnit`.
  ///
  /// - No unit loaded: nothing to do.
  /// - Unit failed: reset it, then wait until systemd garbage-collects it.
  /// - Unit inactive: wait for collection.
  ///
  /// Anything else (active, activating, deactivating, reloading, ...) is an error.
  async fn prepare_name(&self, name: &str) -> Result<(), SystemdTransientError> {
    let proxy = &self.proxy;
    // Subscribe before the first lookup, so a removal in between is not missed.
    let mut removed = proxy.receive_unit_removed().await?;

    let Some(path) = get_unit(proxy, name).await? else {
      return Ok(());
    };

    let unit = UnitProxy::builder(&self.conn)
      .path(path)
      .cache_properties(CacheProperties::No)
      .build()
      .await?;

    if !unit.transient().await? {
      return Err(SystemdTransientError::UnitNotTransient {
        name: name.to_owned(),
      });
    }

    match exit_of(&unit.active_state().await?) {
      Some(UnitExit::Failed) => proxy.reset_failed_unit(name).await?,
      Some(UnitExit::Inactive) => {}
      None => {
        return Err(SystemdTransientError::UnitAlreadyRunning {
          name: name.to_owned(),
        });
      }
    }

    let wait = async {
      while get_unit(proxy, name).await?.is_some() {
        loop {
          let signal = removed
            .next()
            .await
            .ok_or(SystemdTransientError::UnitStreamEnded)?;
          match signal.args() {
            Ok(args) if *args.id() == name => break,
            Ok(_) => {}
            Err(err) => warn!("unparsable UnitRemoved signal: {}", err),
          }
        }
      }
      Ok::<_, SystemdTransientError>(())
    };

    tokio::time::timeout(UNIT_UNLOAD_TIMEOUT, wait)
      .await
      .map_err(|_| SystemdTransientError::UnitStillLoaded {
        name: name.to_owned(),
      })?
  }

  /// Stops a unit by full name (e.g. `foo.slice`) and waits for the job.
  /// A unit systemd does not know counts as stopped.
  pub async fn stop_unit(&self, name: &str) -> Result<(), SystemdTransientError> {
    match self.stop_job(name).await {
      Err(SystemdTransientError::Zbus(zbus::Error::MethodError(err, _, _)))
        if err.as_str() == ERR_NO_SUCH_UNIT =>
      {
        Ok(())
      }
      other => other,
    }
  }

  /// Queues a stop job for `name` and waits for it.
  ///
  /// There is no client-side timeout. systemd enforces the unit's job
  /// timeouts itself (`TimeoutStopUSec`, `JobTimeoutUSec`) and reports an
  /// expiry as a `JobRemoved` result of `timeout`.
  pub(crate) async fn stop_job(&self, name: &str) -> Result<(), SystemdTransientError> {
    let jobs = self.proxy.receive_job_removed().await?;
    let job = self
      .proxy
      .stop_unit(name, JobMode::Replace.as_str())
      .await?;
    wait_for_job(jobs, job.as_ref()).await
  }
}

/// `GetUnit`, with `NoSuchUnit` mapped to `None`. Unlike `LoadUnit`, this
/// does not create a stub unit.
async fn get_unit(
  proxy: &ManagerProxy<'_>,
  name: &str,
) -> Result<Option<OwnedObjectPath>, SystemdTransientError> {
  match proxy.get_unit(name).await {
    Ok(path) => Ok(Some(path)),
    Err(zbus::Error::MethodError(err, _, _)) if err.as_str() == ERR_NO_SUCH_UNIT => Ok(None),
    Err(err) => Err(err.into()),
  }
}

async fn wait_for_job(
  mut stream: JobRemovedStream,
  job_path: ObjectPath<'_>,
) -> Result<(), SystemdTransientError> {
  while let Some(signal) = stream.next().await {
    let args = match signal.args() {
      Ok(args) => args,
      Err(err) => {
        warn!("unparsable JobRemoved signal: {}", err);
        continue;
      }
    };

    if *args.job() != job_path {
      continue;
    }

    return match *args.result() {
      "done" => Ok(()),
      r => Err(SystemdTransientError::JobWaitFailed {
        unit: args.unit().to_string(),
        reason: r.to_owned(),
      }),
    };
  }

  Err(SystemdTransientError::JobStreamEnded)
}

/// A unit as [`SystemdTransientManager::list_units`] reports it.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct ListedUnit {
  pub name: String,
  pub load_state: String,
  pub active_state: String,
  pub sub_state: String,
}

impl ListedUnit {
  /// The state the unit ended in, or `None` while it is active or changing
  /// state.
  pub fn exit(&self) -> Option<UnitExit> {
    exit_of(&self.active_state)
  }
}