soma_infra/
bollard_image_pull.rs

1use futures_util::StreamExt;
2use serde_json::Value;
3use soma_fleet::{HostRecord, TopologyRevision};
4use soma_ops::{MutationSendState, ProgressEvent, Timestamp};
5use tokio_util::sync::CancellationToken;
6
7use crate::{
8    BollardReadClient, ImagePullMutator, ImagePullProgressFrame, ImagePullReceipt,
9    ImagePullRequest, InfraError, MutationFailure, MutationProgressReporter, MutationResult,
10};
11
12#[async_trait::async_trait]
13impl ImagePullMutator for BollardReadClient {
14    async fn pull_image(
15        &self,
16        host: &HostRecord,
17        request: &ImagePullRequest,
18        progress: &dyn MutationProgressReporter,
19        cancellation: &CancellationToken,
20    ) -> MutationResult<ImagePullReceipt> {
21        self.validate_host(host)
22            .map_err(|error| MutationFailure::new(MutationSendState::NotSent, error))?;
23        ensure_not_expired(request.deadline(), cancellation)?;
24        let (from_image, tag) = split_image_reference(request.image());
25        let options = bollard::query_parameters::CreateImageOptions {
26            from_image: Some(from_image),
27            tag,
28            ..Default::default()
29        };
30        let mut stream = self.docker().create_image(Some(options), None, None);
31        let mut receipt = ImagePullReceipt {
32            host: host.id().clone(),
33            topology_revision: TopologyRevision::clone(host.revision()),
34            image: request.image().to_owned(),
35            send_state: MutationSendState::Sent,
36            total_events: 0,
37            progress: Vec::new(),
38            progress_truncated: false,
39            progress_delivery_errors: Vec::new(),
40        };
41        let mut sequence = 0_u64;
42        loop {
43            let remaining = remaining_duration(request.deadline())?;
44            let next = tokio::select! {
45                () = cancellation.cancelled() => {
46                    return Err(MutationFailure::new(
47                        MutationSendState::Unknown,
48                        soma_fleet::FleetError::Cancelled.into(),
49                    ));
50                }
51                () = tokio::time::sleep(remaining) => {
52                    return Err(MutationFailure::new(
53                        MutationSendState::Unknown,
54                        soma_fleet::FleetError::DeadlineExceeded.into(),
55                    ));
56                }
57                next = stream.next() => next,
58            };
59            let Some(item) = next else {
60                break;
61            };
62            let info = item.map_err(|error| {
63                MutationFailure::new(
64                    MutationSendState::Unknown,
65                    InfraError::Docker(error.to_string()),
66                )
67            })?;
68            sequence = sequence.saturating_add(1);
69            let value = serde_json::to_value(&info).map_err(|error| {
70                MutationFailure::new(
71                    MutationSendState::Sent,
72                    InfraError::Parse {
73                        domain: "image-pull",
74                        message: error.to_string(),
75                    },
76                )
77            })?;
78            let frame = progress_frame(sequence, &value);
79            if let Some(error) = frame.error.clone() {
80                receipt.retain_frame(frame);
81                return Err(MutationFailure::new(
82                    MutationSendState::Sent,
83                    InfraError::Docker(error),
84                ));
85            }
86            if let Ok(event) = canonical_progress_event(request, &frame)
87                && let Err(error) = progress.report(&event)
88            {
89                receipt.retain_delivery_error(error);
90            }
91            receipt.retain_frame(frame);
92        }
93        Ok(receipt)
94    }
95}
96
97fn ensure_not_expired(deadline: Timestamp, cancellation: &CancellationToken) -> MutationResult<()> {
98    if cancellation.is_cancelled() {
99        return Err(MutationFailure::new(
100            MutationSendState::NotSent,
101            soma_fleet::FleetError::Cancelled.into(),
102        ));
103    }
104    if Timestamp::now() >= deadline {
105        return Err(MutationFailure::new(
106            MutationSendState::NotSent,
107            soma_fleet::FleetError::DeadlineExceeded.into(),
108        ));
109    }
110    Ok(())
111}
112
113fn remaining_duration(deadline: Timestamp) -> MutationResult<std::time::Duration> {
114    let remaining = deadline
115        .unix_millis()
116        .saturating_sub(Timestamp::now().unix_millis());
117    if remaining <= 0 {
118        Err(MutationFailure::new(
119            MutationSendState::Unknown,
120            soma_fleet::FleetError::DeadlineExceeded.into(),
121        ))
122    } else {
123        Ok(std::time::Duration::from_millis(remaining as u64))
124    }
125}
126
127fn split_image_reference(image: &str) -> (String, Option<String>) {
128    if image.contains('@') {
129        return (image.to_owned(), None);
130    }
131    match image.rsplit_once(':') {
132        Some((repo, tag)) if !tag.contains('/') && !tag.is_empty() => {
133            (repo.to_owned(), Some(tag.to_owned()))
134        }
135        _ => (image.to_owned(), None),
136    }
137}
138
139fn progress_frame(sequence: u64, value: &Value) -> ImagePullProgressFrame {
140    let detail = value
141        .get("progress_detail")
142        .or_else(|| value.get("progressDetail"));
143    ImagePullProgressFrame {
144        sequence,
145        status: text(value, &["status", "Status"]),
146        id: text(value, &["id", "ID", "Id"]),
147        current: detail.and_then(|value| unsigned(value, &["current", "Current"])),
148        total: detail.and_then(|value| unsigned(value, &["total", "Total"])),
149        message: text(value, &["progress", "Progress"]),
150        error: text(value, &["error", "Error"]).or_else(|| {
151            value
152                .get("error_detail")
153                .or_else(|| value.get("errorDetail"))
154                .and_then(|detail| text(detail, &["message", "Message"]))
155        }),
156    }
157}
158
159fn canonical_progress_event(
160    request: &ImagePullRequest,
161    frame: &ImagePullProgressFrame,
162) -> Result<ProgressEvent, soma_ops::ProgressError> {
163    let mut event = ProgressEvent::new(
164        request.operation_id().clone(),
165        request.operation().clone(),
166        frame.sequence,
167        Timestamp::now(),
168        "pull",
169    )?;
170    if let Some(current) = frame.current
171        && frame
172            .total
173            .is_none_or(|total| total > 0 && current <= total)
174    {
175        event = event.with_amount(current, frame.total, Some("bytes"))?;
176    }
177    let message = [
178        frame.status.as_deref(),
179        frame.id.as_deref(),
180        frame.message.as_deref(),
181    ]
182    .into_iter()
183    .flatten()
184    .collect::<Vec<_>>()
185    .join(" | ");
186    if !message.is_empty() {
187        event = event.with_message(bounded_message(&message))?;
188    }
189    Ok(event)
190}
191
192fn bounded_message(value: &str) -> String {
193    value
194        .chars()
195        .map(|character| {
196            if character.is_control() {
197                ' '
198            } else {
199                character
200            }
201        })
202        .take(1_024)
203        .collect()
204}
205
206fn text(value: &Value, names: &[&str]) -> Option<String> {
207    let object = value.as_object()?;
208    names
209        .iter()
210        .find_map(|name| object.get(*name))
211        .and_then(Value::as_str)
212        .map(str::to_owned)
213}
214
215fn unsigned(value: &Value, names: &[&str]) -> Option<u64> {
216    let object = value.as_object()?;
217    names
218        .iter()
219        .find_map(|name| object.get(*name))
220        .and_then(Value::as_u64)
221}
222
223#[cfg(test)]
224#[path = "bollard_image_pull_tests.rs"]
225mod tests;
soma_infra/bollard_image_pull.rs

soma_infra/
bollard_image_pull.rs
