soma_infra/
process_zfs.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use soma_fleet::{CommandExecutor, CommandOutput, CommandRequest, HostRecord};
5use soma_ops::Timestamp;
6use tokio_util::sync::CancellationToken;
7
8use crate::zfs::parse_zfs_table;
9use crate::{
10    InfraError, InfraResult, ZfsDatasetRequest, ZfsInspector, ZfsPoolRequest, ZfsSnapshotRequest,
11    ZfsTable,
12};
13
14const ZFS_OUTPUT_LIMIT: usize = 2 * 1024 * 1024;
15
16/// ZFS inspector backed by a fleet command executor.
17pub struct CommandZfsInspector<E> {
18    executor: Arc<E>,
19}
20
21impl<E> CommandZfsInspector<E> {
22    /// Creates an inspector using the supplied fleet executor.
23    #[must_use]
24    pub fn new(executor: Arc<E>) -> Self {
25        Self { executor }
26    }
27}
28
29#[async_trait]
30impl<E> ZfsInspector for CommandZfsInspector<E>
31where
32    E: CommandExecutor,
33{
34    async fn pools(
35        &self,
36        host: &HostRecord,
37        request: &ZfsPoolRequest,
38        cancellation: &CancellationToken,
39    ) -> InfraResult<ZfsTable> {
40        let mut args = vec!["list".to_owned()];
41        if let Some(pool) = request.pool() {
42            args.push(pool.to_owned());
43        }
44        let output = self
45            .execute(host, "zpool", args, request.deadline(), cancellation)
46            .await?;
47        parse_zfs_table(host, &output, None)
48    }
49
50    async fn datasets(
51        &self,
52        host: &HostRecord,
53        request: &ZfsDatasetRequest,
54        cancellation: &CancellationToken,
55    ) -> InfraResult<ZfsTable> {
56        let mut args = vec!["list".to_owned()];
57        if let Some(dataset_type) = request.dataset_type() {
58            args.extend(["-t".into(), dataset_type.as_arg().into()]);
59        }
60        if request.is_recursive() || request.pool().is_some() {
61            args.push("-r".into());
62        }
63        if let Some(pool) = request.pool() {
64            args.push(pool.to_owned());
65        }
66        let output = self
67            .execute(host, "zfs", args, request.deadline(), cancellation)
68            .await?;
69        parse_zfs_table(host, &output, None)
70    }
71
72    async fn snapshots(
73        &self,
74        host: &HostRecord,
75        request: &ZfsSnapshotRequest,
76        cancellation: &CancellationToken,
77    ) -> InfraResult<ZfsTable> {
78        let mut args = vec!["list".into(), "-t".into(), "snapshot".into()];
79        if let Some(target) = request.dataset().or_else(|| request.pool()) {
80            args.extend(["-r".into(), target.to_owned()]);
81        }
82        let output = self
83            .execute(host, "zfs", args, request.deadline(), cancellation)
84            .await?;
85        parse_zfs_table(host, &output, Some(request.limit()))
86    }
87}
88
89impl<E> CommandZfsInspector<E>
90where
91    E: CommandExecutor,
92{
93    async fn execute(
94        &self,
95        host: &HostRecord,
96        program: &str,
97        args: Vec<String>,
98        deadline: Timestamp,
99        cancellation: &CancellationToken,
100    ) -> InfraResult<String> {
101        let request = CommandRequest::new(program, args, deadline)
102            .map_err(soma_fleet::FleetError::from)?
103            .with_output_limits(ZFS_OUTPUT_LIMIT, ZFS_OUTPUT_LIMIT)
104            .map_err(soma_fleet::FleetError::from)?;
105        let output = self.executor.execute(host, &request, cancellation).await?;
106        validate_output(host, program, &output)?;
107        std::str::from_utf8(output.stdout())
108            .map(str::to_owned)
109            .map_err(|error| InfraError::Parse {
110                domain: "zfs",
111                message: format!("ZFS output is not UTF-8: {error}"),
112            })
113    }
114}
115
116fn validate_output(host: &HostRecord, program: &str, output: &CommandOutput) -> InfraResult<()> {
117    if output.truncated() {
118        return Err(InfraError::InvalidRequest {
119            domain: "zfs",
120            message: format!("{program} output exceeded {ZFS_OUTPUT_LIMIT} bytes"),
121        });
122    }
123    if output.exit_code() != Some(0) {
124        return Err(InfraError::CommandFailed {
125            domain: "zfs",
126            host: host.id().clone(),
127            exit_code: output.exit_code(),
128            stderr: String::from_utf8_lossy(output.stderr()).trim().to_owned(),
129        });
130    }
131    Ok(())
132}
133
134#[cfg(test)]
135#[path = "process_zfs_tests.rs"]
136mod tests;
soma_infra/process_zfs.rs

soma_infra/
process_zfs.rs
