Skip to main content

soma_infra/
bollard_telemetry.rs

1use async_trait::async_trait;
2use bollard::query_parameters::{DataUsageOptions, LogsOptions, StatsOptions};
3use futures_util::StreamExt;
4use soma_fleet::HostRecord;
5use tokio_util::sync::CancellationToken;
6
7use crate::bollard_driver::{BollardReadClient, cancellable};
8use crate::docker_map::parse_error;
9use crate::docker_telemetry_map::{map_container_stats, map_disk_usage};
10use crate::{
11    ContainerLogOptions, ContainerLogs, ContainerStatsSnapshot, DockerDiskUsage, DockerLogStream,
12    DockerTelemetryReader, InfraError, InfraResult,
13};
14
15const MAX_LOG_BYTES: usize = 4 * 1024 * 1024;
16
17#[async_trait]
18impl DockerTelemetryReader for BollardReadClient {
19    async fn disk_usage(
20        &self,
21        host: &HostRecord,
22        cancellation: &CancellationToken,
23    ) -> InfraResult<DockerDiskUsage> {
24        self.validate_host(host)?;
25        let response =
26            cancellable(cancellation, self.docker().df(None::<DataUsageOptions>)).await?;
27        let value =
28            serde_json::to_value(response).map_err(|error| parse_error(error.to_string()))?;
29        map_disk_usage(host, &value)
30    }
31
32    async fn container_logs(
33        &self,
34        host: &HostRecord,
35        container: &str,
36        options: &ContainerLogOptions,
37        cancellation: &CancellationToken,
38    ) -> InfraResult<ContainerLogs> {
39        self.validate_host(host)?;
40        validate_identifier(container)?;
41        let (stdout, stderr) = match options.stream() {
42            DockerLogStream::Stdout => (true, false),
43            DockerLogStream::Stderr => (false, true),
44            DockerLogStream::Both => (true, true),
45        };
46        let query = LogsOptions {
47            follow: false,
48            stdout,
49            stderr,
50            since: to_i32_time("since", options.since_unix_seconds())?,
51            until: to_i32_time("until", options.until_unix_seconds())?,
52            timestamps: false,
53            tail: options.lines().to_string(),
54        };
55        let stream = self.docker().logs(container, Some(query));
56        futures_util::pin_mut!(stream);
57        let mut lines = Vec::new();
58        let mut retained_bytes = 0_usize;
59        let mut truncated = false;
60        'frames: loop {
61            let item = tokio::select! {
62                () = cancellation.cancelled() => {
63                    return Err(soma_fleet::FleetError::Cancelled.into());
64                }
65                item = stream.next() => item,
66            };
67            let Some(item) = item else {
68                break;
69            };
70            let frame = item.map_err(|error| InfraError::Docker(error.to_string()))?;
71            for line in frame
72                .to_string()
73                .lines()
74                .map(|line| line.trim_end_matches(['\r', '\n']))
75                .filter(|line| !line.is_empty())
76                .filter(|line| options.grep().is_none_or(|pattern| line.contains(pattern)))
77            {
78                let next = retained_bytes.saturating_add(line.len());
79                if next > MAX_LOG_BYTES {
80                    truncated = true;
81                    break 'frames;
82                }
83                retained_bytes = next;
84                lines.push(line.to_owned());
85            }
86        }
87        Ok(ContainerLogs {
88            host: host.id().clone(),
89            topology_revision: host.revision().clone(),
90            container: container.to_owned(),
91            lines,
92            truncated,
93        })
94    }
95
96    async fn container_stats(
97        &self,
98        host: &HostRecord,
99        container: &str,
100        cancellation: &CancellationToken,
101    ) -> InfraResult<ContainerStatsSnapshot> {
102        self.validate_host(host)?;
103        validate_identifier(container)?;
104        let stream = self.docker().stats(
105            container,
106            Some(StatsOptions {
107                stream: false,
108                one_shot: true,
109            }),
110        );
111        futures_util::pin_mut!(stream);
112        let item = tokio::select! {
113            () = cancellation.cancelled() => {
114                return Err(soma_fleet::FleetError::Cancelled.into());
115            }
116            item = stream.next() => item,
117        };
118        let stats = item
119            .ok_or_else(|| InfraError::Docker(format!("no stats frame for container {container}")))?
120            .map_err(|error| InfraError::Docker(error.to_string()))?;
121        let value = serde_json::to_value(stats).map_err(|error| parse_error(error.to_string()))?;
122        map_container_stats(host, container, &value)
123    }
124}
125
126fn validate_identifier(value: &str) -> InfraResult<()> {
127    if value.is_empty() || value.len() > 256 || value.chars().any(char::is_control) {
128        Err(InfraError::InvalidRequest {
129            domain: "docker",
130            message: "invalid container identifier".into(),
131        })
132    } else {
133        Ok(())
134    }
135}
136
137fn to_i32_time(name: &'static str, value: Option<i64>) -> InfraResult<i32> {
138    match value {
139        None => Ok(0),
140        Some(value) => i32::try_from(value).map_err(|_| InfraError::InvalidRequest {
141            domain: "docker",
142            message: format!("container log {name} is outside Docker's supported range"),
143        }),
144    }
145}
146
147#[cfg(test)]
148#[path = "bollard_telemetry_tests.rs"]
149mod tests;