Skip to main content

microvms_app/control/
ensure.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Content-addressed build-or-reuse for an image: what `Sandbox::ensure_image` does (#221).
3//!
4//! A harness with many concurrent trials of one task needs one image for them all: the one
5//! already built, the one a sibling trial is building, or — when there is none — one it
6//! builds itself. This module is that composition, from build inputs to a usable image:
7//!
8//! 1. **Name** (IMAGE-6). `<prefix>-<hash12>`, the hash over the artifact's inputs (the
9//!    daemon, the Dockerfile, the build context — [`artifact_content_hash_with_context`])
10//!    plus the base image and the size class, because an image is created on one base at one
11//!    size. Equal inputs name one image; any changed input names a fresh one, which is what
12//!    keeps a stale snapshot from being served under a reused name.
13//! 2. **Artifact** (IMAGE-7). The Dockerfile, the daemon, and the context, zipped with fixed
14//!    dates and modes, so equal inputs are one object at one content-addressed key.
15//! 3. **ARN and upload** (IMAGE-8). The service takes image ARNs, not names, so the ARN is
16//!    built from the caller's account, looked up once per sandbox
17//!    ([`super::services::BuildServices::caller_account`]). The artifact goes to
18//!    `s3://<bucket>/<prefix>/<name>/artifact.zip`, and only when a build is needed.
19//! 4. **Decide** (IMAGE-9, IMAGE-10). After each describe, [`plan`] says what to do: reuse a
20//!    ready image, wait for one building, delete a failed one (or, under `force`, any one)
21//!    and wait for the name to be free, or build.
22//! 5. **Race** (IMAGE-11). Two trials that both find the name free both create; the service
23//!    accepts one and refuses the other. The loser describes again and waits for the winner's
24//!    build, returning it only once it is ready, and marks it reused.
25//!
26//! `crates/model/src/image.rs` is this module as a Stateright model: two concurrent callers against
27//! a platform that moves the image through its states, checked over every interleaving. The
28//! decision table there is [`plan`]'s, row for row.
29//!
30//! # Everything local happens before the first call
31//!
32//! [`prepare`] derives the name, the key, the artifact, and the create request, and runs the
33//! create call's preflight, with zero calls. A request this client refuses costs nothing: not
34//! the account lookup, not a describe, not the upload.
35//!
36//! # Timeouts
37//!
38//! The build wait is the one `build` uses — [`WaitOpts::default`], 45 minutes with a
39//! four-minute stall probe — or the caller's `wait_timeout`. The wait for a deleted name to
40//! free up is [`DEFAULT_DELETE_TIMEOUT`].
41
42use std::collections::BTreeMap;
43use std::time::Duration;
44
45use super::artifact::{BaseImage, artifact_content_hash_with_context, build_artifact_with_context};
46use super::context::BuildContext;
47use super::image::{DEFAULT_BUILD_TIMEOUT, Image, WaitOpts};
48use super::services::BuildServices;
49use super::transport::{Call, paths, send_accepting};
50use super::{ControlPlane, CreateImageRequest, ops};
51use crate::error::{Error, ErrorKind};
52use crate::sizing::SizeClass;
53
54/// How long a deleted image may take to free its name before `ensure_image` gives up.
55///
56/// Five minutes, the bound a harness provider measured for the same wait; a deletion that
57/// has not finished by then is reported rather than waited on for the build's 45.
58pub const DEFAULT_DELETE_TIMEOUT: Duration = Duration::from_secs(300);
59
60/// The longest name the service's `ImageName` admits.
61const MAX_IMAGE_NAME: usize = 64;
62
63/// The hex characters of the identity hash in a name.
64const NAME_HASH_LEN: usize = 12;
65
66/// What a describe of the ensured image found.
67#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
68pub enum Found {
69    /// No image under the name.
70    Absent,
71    /// `CREATING` or `UPDATING`, or a state this client does not know, which is waited on.
72    Building,
73    /// `CREATED` or `UPDATED` (or a tolerated ready spelling): usable.
74    Ready,
75    /// `CREATE_FAILED`, `UPDATE_FAILED`, or `DELETE_FAILED`.
76    Failed,
77    /// `DELETING`.
78    Deleting,
79}
80
81impl Found {
82    /// The reading of a `MicrovmImageState`.
83    pub fn from_state(state: &str) -> Self {
84        if Image::is_ready(state) {
85            Found::Ready
86        } else if Image::is_failed(state) {
87            Found::Failed
88        } else if state == "DELETING" {
89            Found::Deleting
90        } else if state == "DELETED" {
91            Found::Absent
92        } else {
93            Found::Building
94        }
95    }
96}
97
98/// What `ensure_image` does next with what a describe found.
99#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
100pub enum Plan {
101    /// Return the ready image; nothing is uploaded or created.
102    Reuse,
103    /// Wait for the running build to settle.
104    Wait,
105    /// Delete the image, wait until the name is free, then build.
106    Delete,
107    /// A deletion is under way: wait until the name is free, then build.
108    AwaitAbsent,
109    /// Upload the artifact and create the image.
110    Build,
111}
112
113/// The decision `ensure_image` makes after each describe (IMAGE-9, IMAGE-10).
114///
115/// A ready image is reused unless the caller forces a rebuild; a building one is waited on
116/// either way, because the service refuses to delete an image in `CREATING` — a forced
117/// caller deletes what the build settles to. A failed image is deleted whether or not the
118/// caller forced, because the name is content-addressed: the only way to a fresh build under
119/// it is to free it first. `crates/model/src/image.rs` checks this table against two concurrent
120/// callers, and `the_plan_table` holds the two equal.
121pub fn plan(found: Found, force: bool) -> Plan {
122    match (found, force) {
123        (Found::Absent, _) => Plan::Build,
124        (Found::Ready, false) => Plan::Reuse,
125        (Found::Ready, true) | (Found::Failed, _) => Plan::Delete,
126        (Found::Building, _) => Plan::Wait,
127        (Found::Deleting, _) => Plan::AwaitAbsent,
128    }
129}
130
131/// The image name for a prefix and an identity hash: `<prefix>-<hash12>` (IMAGE-6).
132///
133/// The prefix is reduced to what the service's `ImageName` pattern admits
134/// (`[a-zA-Z0-9-_]`), every other character becoming `-`, trimmed of `-` at both ends, and
135/// cut so the whole name fits the pattern's 64 characters. A prefix with nothing left is
136/// refused: the name would be the hash alone, and nothing in the account would say what it
137/// is for.
138pub fn ensured_image_name(prefix: &str, identity_hash: &str) -> Result<String, Error> {
139    let sanitized: String = prefix
140        .chars()
141        .map(|c| {
142            if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
143                c
144            } else {
145                '-'
146            }
147        })
148        .collect();
149    let room = MAX_IMAGE_NAME - 1 - NAME_HASH_LEN;
150    let cut: String = sanitized.trim_matches('-').chars().take(room).collect();
151    let cut = cut.trim_end_matches('-');
152    if cut.is_empty() {
153        return Err(Error::invalid_arg(format!(
154            "the image name prefix {prefix:?} has no character the service's ImageName pattern \
155             admits ([a-zA-Z0-9-_]), so the name would be a bare hash. Pass a prefix that says \
156             what the image is for, such as the task's name."
157        )));
158    }
159    let hash: String = identity_hash.chars().take(NAME_HASH_LEN).collect();
160    Ok(format!("{cut}-{hash}"))
161}
162
163/// The identity an ensured image is named for (IMAGE-6): the artifact's content hash, the
164/// base image's name (what `baseImageArn` names), and the size class's baseline.
165///
166/// Length-prefixed and tagged, so no two field boundaries can collide.
167pub fn image_identity_hash(artifact_hash: &str, base: &BaseImage, size: SizeClass) -> String {
168    pinned_identity_hash(artifact_hash, base, None, size)
169}
170
171/// [`image_identity_hash`] with the base version a build pins (#258): a pinned base builds
172/// on that version whatever the service's default later becomes, so the pin is identity.
173///
174/// `None` is [`image_identity_hash`]'s stream exactly, so an unpinned image keeps the name it
175/// always had; a pin adds one tagged, length-prefixed field after the size.
176pub fn pinned_identity_hash(
177    artifact_hash: &str,
178    base: &BaseImage,
179    base_version: Option<&str>,
180    size: SizeClass,
181) -> String {
182    use sha2::{Digest as _, Sha256};
183
184    let mut hasher = Sha256::new();
185    for field in [
186        b"microvms-ensure-image/1".as_slice(),
187        artifact_hash.as_bytes(),
188        base.name.as_bytes(),
189    ] {
190        hasher.update((field.len() as u64).to_be_bytes());
191        hasher.update(field);
192    }
193    hasher.update(size.baseline_mib().to_be_bytes());
194    if let Some(version) = base_version {
195        for field in [b"baseImageVersion".as_slice(), version.as_bytes()] {
196            hasher.update((field.len() as u64).to_be_bytes());
197            hasher.update(field);
198        }
199    }
200    const_hex::encode(hasher.finalize())
201}
202
203/// The artifact's S3 key: `<prefix>/<name>/artifact.zip`, with the prefix's empty segments
204/// dropped so no key starts with or doubles a `/` (IMAGE-8).
205pub fn artifact_key(key_prefix: Option<&str>, name: &str) -> String {
206    let prefix: Vec<&str> = key_prefix
207        .unwrap_or("")
208        .split('/')
209        .filter(|segment| !segment.is_empty())
210        .collect();
211    if prefix.is_empty() {
212        format!("{name}/artifact.zip")
213    } else {
214        format!("{}/{name}/artifact.zip", prefix.join("/"))
215    }
216}
217
218/// The bucket and key of an `s3://<bucket>/<key>` URI, refused when it isn't one S3 admits.
219///
220/// For the upload of a build whose artifact URI the caller chose ([`crate::sandbox::Sandbox::
221/// upload_artifact_for`]); an ensure derives its own.
222pub fn s3_location(uri: &str) -> Result<(&str, &str), Error> {
223    let (bucket, key) = uri
224        .strip_prefix("s3://")
225        .and_then(|rest| rest.split_once('/'))
226        .filter(|(_, key)| !key.is_empty())
227        .ok_or_else(|| {
228            Error::invalid_arg(format!(
229                "the artifact URI {uri:?} is not s3://<bucket>/<key>, the only place \
230                 CreateMicrovmImage reads an artifact from."
231            ))
232        })?;
233    require_valid_bucket(bucket)?;
234    Ok((bucket, key))
235}
236
237/// Refuses a bucket name S3 does not admit, before any call.
238fn require_valid_bucket(bucket: &str) -> Result<(), Error> {
239    let chars_ok = bucket
240        .bytes()
241        .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'.' || b == b'-');
242    let ends_ok = bucket
243        .bytes()
244        .next()
245        .zip(bucket.bytes().last())
246        .is_some_and(|(first, last)| first.is_ascii_alphanumeric() && last.is_ascii_alphanumeric());
247    if (3..=63).contains(&bucket.len()) && chars_ok && ends_ok && !bucket.contains("..") {
248        return Ok(());
249    }
250    Err(Error::invalid_arg(format!(
251        "the S3 bucket {bucket:?} is not a bucket name S3 admits: 3 to 63 characters of \
252         lowercase letters, digits, `.` and `-`, starting and ending with a letter or digit. \
253         Pass the bucket the build role can read, in the sandbox's region."
254    )))
255}
256
257/// Refuses a key prefix that could not form a legal S3 key.
258fn require_valid_key_prefix(prefix: Option<&str>, name: &str) -> Result<(), Error> {
259    let Some(prefix) = prefix else {
260        return Ok(());
261    };
262    if prefix.chars().any(char::is_control) {
263        return Err(Error::invalid_arg(format!(
264            "the S3 key prefix {prefix:?} carries a control character, which no S3 key \
265             should; pass a plain path such as `harbor/images`."
266        )));
267    }
268    let key = artifact_key(Some(prefix), name);
269    if key.len() > 1024 {
270        return Err(Error::invalid_arg(format!(
271            "the S3 key prefix makes a {}-byte key, over S3's 1024; shorten it.",
272            key.len()
273        )));
274    }
275    Ok(())
276}
277
278/// Everything `Sandbox::ensure_image` needs.
279#[derive(Clone, Debug)]
280pub struct EnsureImageRequest {
281    /// The stem of the image name; the name is `<prefix>-<hash12>`.
282    pub name_prefix: String,
283    /// The daemon binary's bytes.
284    pub binary: Vec<u8>,
285    /// The Dockerfile, usually `wrap_dockerfile`'s output. Checked by the create call's
286    /// preflight before anything else happens.
287    pub dockerfile: String,
288    /// The files the Dockerfile may `COPY` ([`super::context::from_dir`]), or `None`.
289    pub context: Option<BuildContext>,
290    /// The bucket the artifact is uploaded to, in the sandbox's region.
291    pub s3_bucket: String,
292    /// A key prefix inside the bucket, or `None` for the bucket root.
293    pub s3_key_prefix: Option<String>,
294    /// The build role, which must read the bucket and write logs under
295    /// `/aws/lambda-microvms/`.
296    pub build_role_arn: String,
297    /// The size class the image is created with; part of the name.
298    pub size: SizeClass,
299    /// The base image, or `None` for [`BaseImage::from_dockerfile`] of the Dockerfile.
300    pub base_image: Option<BaseImage>,
301    /// Delete what exists under the name and build afresh.
302    pub force: bool,
303    /// Tags for a created image. Not part of the name: a reused image keeps its own.
304    pub tags: BTreeMap<String, String>,
305    /// The build wait's deadline, or `None` for [`DEFAULT_BUILD_TIMEOUT`].
306    pub wait_timeout: Option<Duration>,
307    /// A project's manifest and lockfile, baked as an environment layer (#74). Part of the
308    /// name: they enter the artifact.
309    pub project_files: Option<super::artifact::ProjectFiles>,
310    /// `baseImageVersion`, or `None` for the service's default. Part of the name when set.
311    pub base_image_version: Option<String>,
312    /// `logging.cloudWatch.logGroup` for a created image. Not part of the name: logging
313    /// doesn't change what is built, and a reused image keeps its own.
314    pub log_group: Option<String>,
315    /// The log-stream prefix for a created image; see [`CreateImageRequest::log_stream`].
316    pub log_stream: Option<String>,
317    /// [`CreateImageRequest::repair_guest_identity`] for a created image. Not part of the
318    /// name, as `build --reuse`'s name never held it: a reused image keeps its own.
319    pub repair_guest_identity: bool,
320    /// [`CreateImageRequest::inherit_workdir`] for a created image.
321    pub inherit_workdir: bool,
322    /// The run-family hook timeouts for a created image.
323    pub run_hook_timeout: crate::hooks::RunHookTimeout,
324    /// The build-family hook timeouts for a created image.
325    pub build_hook_timeout: crate::hooks::BuildHookTimeout,
326}
327
328impl EnsureImageRequest {
329    /// A request with the defaults: no context, no key prefix, the default size class, the
330    /// base derived from the Dockerfile, no force, no tags, the default wait, and the create
331    /// request's own defaults for everything a created image carries.
332    pub fn new(
333        name_prefix: impl Into<String>,
334        binary: Vec<u8>,
335        dockerfile: impl Into<String>,
336        s3_bucket: impl Into<String>,
337        build_role_arn: impl Into<String>,
338    ) -> Self {
339        // The create request's defaults, read off it rather than restated, so a created
340        // image carries what `build_image` would have given it.
341        let defaults = CreateImageRequest::new("", Vec::new(), "", "");
342        Self {
343            name_prefix: name_prefix.into(),
344            binary,
345            dockerfile: dockerfile.into(),
346            context: None,
347            s3_bucket: s3_bucket.into(),
348            s3_key_prefix: None,
349            build_role_arn: build_role_arn.into(),
350            size: SizeClass::DEFAULT,
351            base_image: None,
352            force: false,
353            tags: BTreeMap::new(),
354            wait_timeout: None,
355            project_files: None,
356            base_image_version: None,
357            log_group: None,
358            log_stream: None,
359            repair_guest_identity: defaults.repair_guest_identity,
360            inherit_workdir: defaults.inherit_workdir,
361            run_hook_timeout: defaults.run_hook_timeout,
362            build_hook_timeout: defaults.build_hook_timeout,
363        }
364    }
365
366    /// The ensure request for a create request's inputs (#258): its name becomes the prefix
367    /// and every field it carries a created image carries too. `dockerfile` is the text the
368    /// build takes, which a caller with no Dockerfile of its own gets from
369    /// [`ControlPlane::dockerfile_for`]. The artifact URI and token label are dropped: the
370    /// ensure derives both from the name it computes.
371    pub fn from_create(
372        create: CreateImageRequest,
373        dockerfile: String,
374        s3_bucket: impl Into<String>,
375    ) -> Self {
376        let mut request = Self::new(
377            create.name,
378            create.binary,
379            dockerfile,
380            s3_bucket,
381            create.build_role_arn,
382        );
383        request.size = create.size;
384        request.base_image = Some(create.base_image);
385        request.tags = create.tags;
386        request.project_files = create.project_files;
387        request.base_image_version = create.base_image_version;
388        request.log_group = create.log_group;
389        request.log_stream = create.log_stream;
390        request.repair_guest_identity = create.repair_guest_identity;
391        request.inherit_workdir = create.inherit_workdir;
392        request.run_hook_timeout = create.run_hook_timeout;
393        request.build_hook_timeout = create.build_hook_timeout;
394        request
395    }
396}
397
398/// What `Sandbox::ensure_image` returns.
399#[derive(Clone, Debug)]
400pub struct EnsuredImage {
401    /// The ready image.
402    pub image: Image,
403    /// True when this call's own create did not build the image: it was ready, or a build
404    /// already running was waited out, or a sibling won the create race.
405    pub reused: bool,
406    /// Where the artifact is, whether or not this call uploaded it.
407    pub artifact_uri: String,
408    /// Whether this call uploaded the artifact.
409    pub uploaded: bool,
410    /// What reading the build context skipped, one line each ([`BuildContext::warnings`]).
411    pub warnings: Vec<String>,
412}
413
414/// Everything [`prepare`] decided locally: the name, the key, the artifact, and the create
415/// request, with the create's preflight already passed.
416#[derive(Clone, Debug)]
417pub struct Prepared {
418    /// `<prefix>-<hash12>`.
419    pub name: String,
420    pub bucket: String,
421    pub key: String,
422    /// The zip to upload when a build is needed.
423    pub artifact: Vec<u8>,
424    /// The create request, named and pointed at the artifact.
425    pub create: CreateImageRequest,
426    pub force: bool,
427    pub wait_timeout: Option<Duration>,
428    pub warnings: Vec<String>,
429}
430
431impl Prepared {
432    /// `s3://<bucket>/<key>`.
433    pub fn artifact_uri(&self) -> String {
434        format!("s3://{}/{}", self.bucket, self.key)
435    }
436}
437
438/// The local half of `ensure_image`: every name, key, and artifact decision and every guard,
439/// with zero calls.
440pub fn prepare(control: &ControlPlane, request: EnsureImageRequest) -> Result<Prepared, Error> {
441    require_valid_bucket(&request.s3_bucket)?;
442    let base = match request.base_image {
443        Some(base) => base,
444        None => BaseImage::from_dockerfile(&request.dockerfile)?,
445    };
446    let artifact_hash = artifact_content_hash_with_context(
447        &request.binary,
448        &request.dockerfile,
449        request.project_files.as_ref(),
450        request.context.as_ref(),
451    );
452    let name = ensured_image_name(
453        &request.name_prefix,
454        &pinned_identity_hash(
455            &artifact_hash,
456            &base,
457            request.base_image_version.as_deref(),
458            request.size,
459        ),
460    )?;
461    require_valid_key_prefix(request.s3_key_prefix.as_deref(), &name)?;
462    let key = artifact_key(request.s3_key_prefix.as_deref(), &name);
463
464    let mut create = CreateImageRequest::new(
465        name.clone(),
466        request.binary,
467        format!("s3://{}/{key}", request.s3_bucket),
468        request.build_role_arn,
469    );
470    create.base_image = base;
471    create.dockerfile = Some(request.dockerfile);
472    create.size = request.size;
473    create.tags = request.tags;
474    create.token_scope = Some(name.clone());
475    create.project_files = request.project_files;
476    create.base_image_version = request.base_image_version;
477    create.log_group = request.log_group;
478    create.log_stream = request.log_stream;
479    create.repair_guest_identity = request.repair_guest_identity;
480    create.inherit_workdir = request.inherit_workdir;
481    create.run_hook_timeout = request.run_hook_timeout;
482    create.build_hook_timeout = request.build_hook_timeout;
483    control.preflight(&create)?;
484
485    let artifact = build_artifact_with_context(
486        &create.binary,
487        create.dockerfile.as_deref().unwrap_or_default(),
488        create.project_files.as_ref(),
489        request.context.as_ref(),
490    )?;
491    Ok(Prepared {
492        name,
493        bucket: request.s3_bucket,
494        key,
495        artifact,
496        create,
497        force: request.force,
498        wait_timeout: request.wait_timeout,
499        warnings: request
500            .context
501            .map(|context| context.warnings().to_vec())
502            .unwrap_or_default(),
503    })
504}
505
506impl ControlPlane {
507    /// The image under `identifier`, or `None` when the service answers that there is none.
508    ///
509    /// `GetMicrovmImage`'s 404 is an answer here rather than an error: "no image under this
510    /// name" is the state `ensure_image` builds from. Every other failure is an error as
511    /// usual, and throttles are retried.
512    pub async fn describe_image(
513        &self,
514        identifier: &str,
515    ) -> Result<Option<ops::GetMicrovmImageResponseWire>, Error> {
516        super::require_valid_identifier("imageIdentifier", identifier)?;
517        let call = Call::get("GetMicrovmImage", paths::microvm_image(identifier));
518        let reply = send_accepting(self.transport(), call, &[404]).await?;
519        if reply.status == 404 {
520            return Ok(None);
521        }
522        Ok(Some(reply.json("GetMicrovmImage")?))
523    }
524
525    /// Polls until `identifier` is no longer building, answering what it settled to.
526    async fn await_settled(&self, identifier: &str, opts: WaitOpts) -> Result<Found, Error> {
527        let started = self.clock().elapsed();
528        loop {
529            let found = self
530                .describe_image(identifier)
531                .await?
532                .map_or(Found::Absent, |image| Found::from_state(&image.state));
533            if found != Found::Building {
534                return Ok(found);
535            }
536            let elapsed = self.clock().elapsed().saturating_sub(started);
537            if elapsed >= opts.timeout {
538                return Err(super::timed_out(
539                    &format!("image {identifier} did not finish building"),
540                    elapsed,
541                ));
542            }
543            self.clock().sleep(opts.poll_interval).await;
544        }
545    }
546
547    /// Deletes `identifier` and polls until its name is free (IMAGE-10).
548    ///
549    /// The delete is [`ControlPlane::delete_image`]'s, versions first. The poll ends when the
550    /// describe answers 404, or when the image under the name is no longer the one deleted —
551    /// a sibling rebuilt it, and the create that follows will join that build — or with
552    /// [`ErrorKind::Platform`] for `DELETE_FAILED` and a timeout at
553    /// [`DEFAULT_DELETE_TIMEOUT`].
554    async fn delete_and_await_absent(
555        &self,
556        identifier: &str,
557        deleted_state: &str,
558        poll_interval: Duration,
559    ) -> Result<(), Error> {
560        self.delete_image(identifier, 3, poll_interval).await;
561        self.await_absent(identifier, Some(deleted_state), poll_interval)
562            .await
563    }
564
565    /// Polls until the name is free, treating `deleted_state` — the state a delete was just
566    /// issued against — as a deletion the describe has not caught up with yet.
567    async fn await_absent(
568        &self,
569        identifier: &str,
570        deleted_state: Option<&str>,
571        poll_interval: Duration,
572    ) -> Result<(), Error> {
573        let started = self.clock().elapsed();
574        loop {
575            let Some(image) = self.describe_image(identifier).await? else {
576                return Ok(());
577            };
578            if image.state == "DELETE_FAILED" {
579                return Err(Error::new(
580                    ErrorKind::Platform,
581                    format!(
582                        "image {identifier} reports DELETE_FAILED, so its name cannot be \
583                         freed for a rebuild. Delete it by hand (`aws lambda-microvms \
584                         delete-microvm-image`), or pass a different name prefix."
585                    ),
586                ));
587            }
588            let pending = image.state == "DELETING" || Some(image.state.as_str()) == deleted_state;
589            if !pending {
590                return Ok(());
591            }
592            let elapsed = self.clock().elapsed().saturating_sub(started);
593            if elapsed >= DEFAULT_DELETE_TIMEOUT {
594                return Err(super::timed_out(
595                    &format!("image {identifier} was still present after its deletion"),
596                    elapsed,
597                ));
598            }
599            self.clock().sleep(poll_interval).await;
600        }
601    }
602
603    /// Whether the name is free or being freed: the image a wait was on has gone.
604    async fn vanished(&self, identifier: &str) -> Result<bool, Error> {
605        Ok(match self.describe_image(identifier).await? {
606            None => true,
607            Some(image) => Found::from_state(&image.state) == Found::Deleting,
608        })
609    }
610}
611
612fn image_from_wire(image: ops::GetMicrovmImageResponseWire, size: SizeClass) -> Image {
613    Image {
614        identifier: image.image_arn,
615        name: image.name,
616        version: image.latest_active_image_version.unwrap_or_default(),
617        state: image.state,
618        size,
619        log_group: None,
620        log_stream: None,
621    }
622}
623
624/// Where one `ensure_image` call is, across its describes.
625#[derive(Default)]
626struct Progress {
627    uploaded: bool,
628    /// This call deleted what it found once; a forced rebuild is one-shot.
629    deleted: bool,
630    /// This call went back to the describe after the image it waited on disappeared.
631    redescribed: bool,
632}
633
634/// The remote half of `ensure_image`: describe, decide, and build, reuse, wait, or join.
635pub(crate) async fn ensure(
636    control: &ControlPlane,
637    services: &dyn BuildServices,
638    account: &str,
639    prepared: Prepared,
640) -> Result<EnsuredImage, Error> {
641    let arn = format!(
642        "arn:aws:lambda:{}:{account}:microvm-image:{}",
643        control.region().as_str(),
644        prepared.name
645    );
646    let size = prepared.create.size;
647    let wait = WaitOpts {
648        timeout: prepared.wait_timeout.unwrap_or(DEFAULT_BUILD_TIMEOUT),
649        ..WaitOpts::default()
650    };
651    let done = |image: Image, reused: bool, uploaded: bool| EnsuredImage {
652        image,
653        reused,
654        artifact_uri: prepared.artifact_uri(),
655        uploaded,
656        warnings: prepared.warnings.clone(),
657    };
658    let mut progress = Progress::default();
659
660    loop {
661        let described = control.describe_image(&arn).await?;
662        let found = described
663            .as_ref()
664            .map_or(Found::Absent, |image| Found::from_state(&image.state));
665        let force = prepared.force && !progress.deleted;
666        match plan(found, force) {
667            Plan::Reuse => {
668                let image = described.expect("a ready image was described");
669                return Ok(done(image_from_wire(image, size), true, progress.uploaded));
670            }
671            Plan::Wait if force => {
672                // The service refuses to delete a build in progress: let it settle, then
673                // delete whatever it settled to.
674                let settled = control.await_settled(&arn, wait).await?;
675                if settled != Found::Absent {
676                    let state = control
677                        .describe_image(&arn)
678                        .await?
679                        .map(|image| image.state)
680                        .unwrap_or_default();
681                    control
682                        .delete_and_await_absent(&arn, &state, wait.poll_interval)
683                        .await?;
684                }
685                progress.deleted = true;
686            }
687            Plan::Wait => match control.wait_for_image(&arn, size, wait).await {
688                Ok(image) => return Ok(done(image, true, progress.uploaded)),
689                Err(error) => {
690                    if !progress.redescribed && control.vanished(&arn).await? {
691                        progress.redescribed = true;
692                        continue;
693                    }
694                    return Err(error);
695                }
696            },
697            Plan::Delete => {
698                let state = described.map(|image| image.state).unwrap_or_default();
699                control
700                    .delete_and_await_absent(&arn, &state, wait.poll_interval)
701                    .await?;
702                progress.deleted = true;
703            }
704            Plan::AwaitAbsent => control.await_absent(&arn, None, wait.poll_interval).await?,
705            Plan::Build => {}
706        }
707
708        // The name is free as far as this call knows: upload, create, and wait.
709        if !progress.uploaded {
710            services
711                .put_object(&prepared.bucket, &prepared.key, prepared.artifact.clone())
712                .await?;
713            progress.uploaded = true;
714        }
715        match control.create_image(prepared.create.clone()).await {
716            Ok(created) => match control
717                .wait_for_image(&created.identifier, size, wait)
718                .await
719            {
720                Ok(mut built) => {
721                    // The logging config survives the wait from the create's own record, as
722                    // in `Sandbox::build_image`: the readback cannot carry it.
723                    built.log_group = created.log_group;
724                    built.log_stream = created.log_stream;
725                    return Ok(done(built, false, true));
726                }
727                Err(error) => {
728                    if !progress.redescribed && control.vanished(&arn).await? {
729                        progress.redescribed = true;
730                        continue;
731                    }
732                    return Err(error);
733                }
734            },
735            // IMAGE-11: a refused create is a race lost when something now holds the name.
736            Err(refused) if refused.kind() == ErrorKind::Platform => {
737                match control.describe_image(&arn).await? {
738                    Some(image) if Found::from_state(&image.state) == Found::Ready => {
739                        return Ok(done(image_from_wire(image, size), true, true));
740                    }
741                    Some(image)
742                        if matches!(
743                            Found::from_state(&image.state),
744                            Found::Building | Found::Failed
745                        ) =>
746                    {
747                        // Waits for the winner's build; a failed build answers its own
748                        // diagnosis rather than this call's refused create.
749                        match control.wait_for_image(&arn, size, wait).await {
750                            Ok(image) => return Ok(done(image, true, true)),
751                            Err(error) => {
752                                if !progress.redescribed && control.vanished(&arn).await? {
753                                    progress.redescribed = true;
754                                    continue;
755                                }
756                                return Err(error);
757                            }
758                        }
759                    }
760                    _ if !progress.redescribed => {
761                        progress.redescribed = true;
762                        continue;
763                    }
764                    _ => return Err(refused),
765                }
766            }
767            Err(error) => return Err(error),
768        }
769    }
770}
771
772#[cfg(test)]
773mod tests {
774    use std::sync::atomic::{AtomicUsize, Ordering};
775    use std::sync::{Arc, Mutex};
776
777    use super::*;
778    use crate::control::fake::{self, Answer, FakeControlPlane, TestClock};
779    use crate::region::Region;
780    use crate::sandbox::Sandbox;
781    use crate::testing::aarch64_daemon;
782
783    const ACCOUNT: &str = "123456789012";
784    const BUCKET: &str = "artifact-bucket";
785    const ROLE: &str = "arn:aws:iam::123456789012:role/build";
786
787    /// STS and S3, recorded.
788    #[derive(Default)]
789    struct FakeServices {
790        account_calls: AtomicUsize,
791        puts: Mutex<Vec<(String, String, Vec<u8>)>>,
792    }
793
794    impl BuildServices for FakeServices {
795        fn caller_account(&self) -> futures_util::future::BoxFuture<'_, Result<String, Error>> {
796            self.account_calls.fetch_add(1, Ordering::SeqCst);
797            Box::pin(async { Ok(ACCOUNT.to_string()) })
798        }
799
800        fn put_object<'a>(
801            &'a self,
802            bucket: &'a str,
803            key: &'a str,
804            bytes: Vec<u8>,
805        ) -> futures_util::future::BoxFuture<'a, Result<(), Error>> {
806            self.puts.lock().expect("not poisoned").push((
807                bucket.to_string(),
808                key.to_string(),
809                bytes,
810            ));
811            Box::pin(async { Ok(()) })
812        }
813    }
814
815    fn sandbox(fake: &Arc<FakeControlPlane>, services: &Arc<FakeServices>) -> Sandbox {
816        let plane = crate::testing::control_plane(
817            fake.clone(),
818            Region::UsEast1,
819            Arc::new(TestClock::new()),
820        );
821        Sandbox::with_control_plane(plane).with_build_services(services.clone())
822    }
823
824    fn dockerfile() -> String {
825        super::super::artifact::wrap_dockerfile(
826            "FROM python:3.12-slim\nWORKDIR /app\n",
827            &super::super::artifact::WrapOptions::default(),
828        )
829        .expect("wraps")
830    }
831
832    fn request() -> EnsureImageRequest {
833        let mut request = EnsureImageRequest::new(
834            "task",
835            aarch64_daemon(b"daemon"),
836            dockerfile(),
837            BUCKET,
838            ROLE,
839        );
840        request.s3_key_prefix = Some("harbor/images".to_string());
841        request
842    }
843
844    /// The name `request()` derives, computed from the public pieces.
845    fn expected_name() -> String {
846        let request = request();
847        let base = BaseImage::from_dockerfile(&request.dockerfile).expect("a FROM");
848        let hash = super::super::artifact::artifact_content_hash_with_context(
849            &request.binary,
850            &request.dockerfile,
851            None,
852            None,
853        );
854        ensured_image_name("task", &image_identity_hash(&hash, &base, request.size))
855            .expect("a legal name")
856    }
857
858    /// `GetMicrovmImageOutput` for a usable image.
859    fn ready(name: &str) -> String {
860        format!(
861            r#"{{"imageArn": "arn:aws:lambda:us-east-1:{ACCOUNT}:microvm-image:{name}",
862                "name": "{name}", "state": "CREATED", "latestActiveImageVersion": "1",
863                "createdAt": 1754524800, "updatedAt": 1754528400, "tags": {{}}}}"#
864        )
865    }
866
867    fn state(name: &str, state: &str) -> Answer {
868        Answer::ok(fake::get_image_response(name, state))
869    }
870
871    fn absent() -> Answer {
872        Answer::failure(404, "Image not found")
873    }
874
875    fn uploads(services: &FakeServices) -> Vec<(String, String)> {
876        services
877            .puts
878            .lock()
879            .expect("not poisoned")
880            .iter()
881            .map(|(bucket, key, _)| (bucket.clone(), key.clone()))
882            .collect()
883    }
884
885    /// **The table `crates/model/src/image.rs` specifies**, row for row (`the_plan_table` there).
886    #[test]
887    fn the_plan_table() {
888        for (found, force, expected) in [
889            (Found::Absent, false, Plan::Build),
890            (Found::Absent, true, Plan::Build),
891            (Found::Ready, false, Plan::Reuse),
892            (Found::Ready, true, Plan::Delete),
893            (Found::Building, false, Plan::Wait),
894            (Found::Building, true, Plan::Wait),
895            (Found::Failed, false, Plan::Delete),
896            (Found::Failed, true, Plan::Delete),
897            (Found::Deleting, false, Plan::AwaitAbsent),
898            (Found::Deleting, true, Plan::AwaitAbsent),
899        ] {
900            assert_eq!(plan(found, force), expected, "{found:?} force={force}");
901        }
902        for (spelling, found) in [
903            ("CREATING", Found::Building),
904            ("UPDATING", Found::Building),
905            ("CREATED", Found::Ready),
906            ("UPDATED", Found::Ready),
907            ("ACTIVE", Found::Ready),
908            ("CREATE_FAILED", Found::Failed),
909            ("UPDATE_FAILED", Found::Failed),
910            ("DELETE_FAILED", Found::Failed),
911            ("DELETING", Found::Deleting),
912            ("DELETED", Found::Absent),
913        ] {
914            assert_eq!(Found::from_state(spelling), found, "{spelling}");
915        }
916    }
917
918    /// **IMAGE-6, the name.** `<prefix>-<hash12>`, the prefix reduced to the characters the
919    /// service's `ImageName` admits and cut so the whole fits its 64; a prefix with nothing
920    /// left is refused.
921    #[test]
922    fn the_name_is_the_sanitized_prefix_and_twelve_hex() {
923        let hash = "0123456789abcdef".repeat(4);
924        assert_eq!(
925            ensured_image_name("harbor/task:1", &hash).expect("legal"),
926            "harbor-task-1-0123456789ab"
927        );
928        let long = ensured_image_name(&"p".repeat(100), &hash).expect("legal");
929        assert_eq!(long.len(), 64, "IMAGE-6: {long}");
930        assert!(long.ends_with("-0123456789ab"));
931        super::super::require_valid_image_name(&long).expect("the service's pattern");
932        for empty in ["", "///", "   "] {
933            let error = ensured_image_name(empty, &hash).expect_err("nothing left");
934            assert_eq!(error.kind(), ErrorKind::InvalidArg);
935        }
936    }
937
938    /// **IMAGE-6, the identity.** The base image and the size class are part of the name,
939    /// because an image is created at one size on one base: two tasks with one Dockerfile and
940    /// different sizes need two images.
941    #[test]
942    fn the_identity_covers_the_base_and_the_size() {
943        let hash = "a".repeat(64);
944        let base = BaseImage::al2023();
945        let identity = image_identity_hash(&hash, &base, SizeClass::DEFAULT);
946        assert_eq!(
947            identity,
948            image_identity_hash(&hash, &base, SizeClass::DEFAULT)
949        );
950        assert_eq!(identity.len(), 64);
951        let other_size = SizeClass::from_baseline_mib(4096).expect("a class");
952        assert_ne!(
953            identity,
954            image_identity_hash(&hash, &base, other_size),
955            "IMAGE-6: size"
956        );
957        let other_base = BaseImage {
958            name: "al2023-2".to_string(),
959            ..BaseImage::al2023()
960        };
961        assert_ne!(
962            identity,
963            image_identity_hash(&hash, &other_base, SizeClass::DEFAULT)
964        );
965        assert_ne!(
966            identity,
967            image_identity_hash(&"b".repeat(64), &base, SizeClass::DEFAULT)
968        );
969    }
970
971    /// **IMAGE-6 with a pinned base (#258).** An unpinned identity is the historical one
972    /// exactly, so every ensured image keeps its name; a pin is identity, and two pins differ.
973    #[test]
974    fn a_pinned_base_version_is_part_of_the_identity() {
975        let hash = "a".repeat(64);
976        let base = BaseImage::al2023();
977        let unpinned = image_identity_hash(&hash, &base, SizeClass::DEFAULT);
978        assert_eq!(
979            pinned_identity_hash(&hash, &base, None, SizeClass::DEFAULT),
980            unpinned,
981            "no pin, the name every ensured image already has"
982        );
983        let one = pinned_identity_hash(&hash, &base, Some("1"), SizeClass::DEFAULT);
984        assert_ne!(one, unpinned);
985        assert_ne!(
986            one,
987            pinned_identity_hash(&hash, &base, Some("0"), SizeClass::DEFAULT)
988        );
989    }
990
991    /// **The create fields an ensure carries (#258).** Project files and a pinned base name a
992    /// different image; logging, identity repair, `inherit_workdir` and the hook timeouts
993    /// don't, and all of them reach the create request a build sends.
994    #[test]
995    fn a_created_image_carries_every_field_of_the_request() {
996        use crate::control::artifact::{Ecosystem, ProjectFiles};
997
998        let plane = crate::testing::control_plane(
999            Arc::new(FakeControlPlane::new()),
1000            Region::UsEast1,
1001            Arc::new(TestClock::new()),
1002        );
1003        let plain = prepare(&plane, request()).expect("prepares");
1004
1005        let mut with_project = request();
1006        with_project.project_files = Some(ProjectFiles {
1007            ecosystem: Ecosystem::Uv,
1008            manifest: b"[project]\nname = \"demo\"\n".to_vec(),
1009            lockfile: b"version = 1\n".to_vec(),
1010        });
1011        // The wrapped Dockerfile doesn't COPY the pair, which core refuses (#74), so the
1012        // project's own default stanza stands in.
1013        with_project.dockerfile = crate::control::artifact::default_dockerfile(
1014            crate::control::DEFAULT_AGENT_PORT,
1015            None,
1016            &BaseImage::al2023(),
1017            Some(Ecosystem::Uv),
1018        );
1019        with_project.base_image = Some(BaseImage::al2023());
1020        let mut without_project = with_project.clone();
1021        without_project.project_files = None;
1022        let project = prepare(&plane, with_project).expect("prepares");
1023        assert!(project.create.project_files.is_some());
1024        assert_ne!(
1025            project.name,
1026            prepare(&plane, without_project).expect("prepares").name,
1027            "the project files enter the artifact, so they name a different image"
1028        );
1029
1030        let mut pinned = request();
1031        pinned.base_image_version = Some("1".to_string());
1032        let pinned = prepare(&plane, pinned).expect("prepares");
1033        assert_ne!(pinned.name, plain.name, "a pin names a different image");
1034        assert_eq!(pinned.create.base_image_version.as_deref(), Some("1"));
1035
1036        let mut configured = request();
1037        configured.log_group = Some("/team/builds".to_string());
1038        configured.log_stream = Some("task".to_string());
1039        configured.repair_guest_identity = true;
1040        configured.inherit_workdir = true;
1041        configured.run_hook_timeout =
1042            crate::hooks::RunHookTimeout::try_new(7).expect("a legal timeout");
1043        let configured = prepare(&plane, configured).expect("prepares");
1044        assert_eq!(configured.name, plain.name, "configuration isn't identity");
1045        assert_eq!(configured.create.log_group.as_deref(), Some("/team/builds"));
1046        assert_eq!(configured.create.log_stream.as_deref(), Some("task"));
1047        assert!(configured.create.repair_guest_identity);
1048        assert!(configured.create.inherit_workdir);
1049        assert_eq!(configured.create.run_hook_timeout.as_secs(), 7);
1050    }
1051
1052    /// `from_create` hands an ensure every field a create request carries, and names the
1053    /// image the way an ensure of the same inputs would.
1054    #[test]
1055    fn an_ensure_from_a_create_request_keeps_its_fields() {
1056        let mut create = CreateImageRequest::new("task", aarch64_daemon(b"daemon"), "", ROLE);
1057        create.dockerfile = Some(dockerfile());
1058        create.base_image = BaseImage::from_dockerfile(&dockerfile()).expect("a FROM");
1059        create.base_image_version = Some("1".to_string());
1060        create.log_group = Some("/g".to_string());
1061        create.repair_guest_identity = true;
1062        create.tags.insert("team".to_string(), "x".to_string());
1063        let from = EnsureImageRequest::from_create(create, dockerfile(), BUCKET);
1064        assert_eq!(from.name_prefix, "task");
1065        assert_eq!(from.s3_bucket, BUCKET);
1066        assert_eq!(from.base_image_version.as_deref(), Some("1"));
1067        assert_eq!(from.log_group.as_deref(), Some("/g"));
1068        assert!(from.repair_guest_identity);
1069        assert_eq!(from.tags.get("team").map(String::as_str), Some("x"));
1070
1071        let mut direct = request();
1072        direct.base_image_version = Some("1".to_string());
1073        direct.s3_key_prefix = None;
1074        let plane = crate::testing::control_plane(
1075            Arc::new(FakeControlPlane::new()),
1076            Region::UsEast1,
1077            Arc::new(TestClock::new()),
1078        );
1079        assert_eq!(
1080            prepare(&plane, from).expect("prepares").name,
1081            prepare(&plane, direct).expect("prepares").name
1082        );
1083    }
1084
1085    /// **A caller-named build uploads through the build services (#258).** The artifact goes to
1086    /// the bucket and key its URI names, once, with the account unasked; a URI that isn't one
1087    /// puts nothing.
1088    #[tokio::test]
1089    async fn a_named_builds_artifact_is_put_where_its_uri_says() {
1090        let fake = Arc::new(FakeControlPlane::new());
1091        let services = Arc::new(FakeServices::default());
1092        let mut sandbox = sandbox(&fake, &services);
1093        let mut create = CreateImageRequest::new(
1094            "named",
1095            aarch64_daemon(b"daemon"),
1096            "s3://artifact-bucket/builds/named.zip",
1097            ROLE,
1098        );
1099        create.dockerfile = Some(dockerfile());
1100        create.base_image = BaseImage::from_dockerfile(&dockerfile()).expect("a FROM");
1101        sandbox.upload_artifact_for(&create).await.expect("uploads");
1102        let puts = services.puts.lock().expect("not poisoned").clone();
1103        assert_eq!(puts.len(), 1);
1104        assert_eq!(
1105            (puts[0].0.as_str(), puts[0].1.as_str()),
1106            ("artifact-bucket", "builds/named.zip")
1107        );
1108        assert_eq!(
1109            puts[0].2,
1110            sandbox
1111                .build_artifact_for(&create)
1112                .expect("the same artifact")
1113        );
1114        assert_eq!(services.account_calls.load(Ordering::SeqCst), 0);
1115        assert!(fake.calls().is_empty(), "no control-plane call");
1116
1117        create.code_artifact_uri = "artifact-bucket/named.zip".to_string();
1118        let error = sandbox
1119            .upload_artifact_for(&create)
1120            .await
1121            .expect_err("not an S3 URI");
1122        assert_eq!(error.kind(), ErrorKind::InvalidArg, "{error}");
1123        assert_eq!(uploads(&services).len(), 1, "nothing more was put");
1124    }
1125
1126    /// An artifact URI is `s3://<bucket>/<key>` with a bucket S3 admits, or it is refused.
1127    #[test]
1128    fn an_artifact_uri_names_a_bucket_and_a_key() {
1129        assert_eq!(
1130            s3_location("s3://artifact-bucket/a/b.zip").expect("an S3 URI"),
1131            ("artifact-bucket", "a/b.zip")
1132        );
1133        for uri in [
1134            "https://artifact-bucket/a.zip",
1135            "s3://artifact-bucket",
1136            "s3://artifact-bucket/",
1137            "s3://Not_A_Bucket/a.zip",
1138        ] {
1139            let error = s3_location(uri).expect_err(uri);
1140            assert_eq!(error.kind(), ErrorKind::InvalidArg, "{uri}: {error}");
1141        }
1142    }
1143
1144    /// **IMAGE-8, the key.** `<prefix>/<name>/artifact.zip`, with the prefix's slashes
1145    /// trimmed so no key carries an empty segment.
1146    #[test]
1147    fn the_artifact_key_is_content_addressed() {
1148        assert_eq!(
1149            artifact_key(Some("harbor/images"), "task-0123456789ab"),
1150            "harbor/images/task-0123456789ab/artifact.zip"
1151        );
1152        assert_eq!(artifact_key(Some("/harbor/"), "n"), "harbor/n/artifact.zip");
1153        assert_eq!(artifact_key(None, "n"), "n/artifact.zip");
1154        assert_eq!(artifact_key(Some(""), "n"), "n/artifact.zip");
1155    }
1156
1157    /// **IMAGE-9, reuse.** A ready image is returned as it is: one account lookup, one
1158    /// describe, no upload, no create.
1159    #[tokio::test]
1160    async fn a_ready_image_is_reused_with_no_upload_and_no_create() {
1161        let name = expected_name();
1162        let fake = Arc::new(FakeControlPlane::new());
1163        fake.answer("GetMicrovmImage", Answer::ok(ready(&name)));
1164        let services = Arc::new(FakeServices::default());
1165        let ensured = sandbox(&fake, &services)
1166            .ensure_image(request())
1167            .await
1168            .expect("reused");
1169        assert!(ensured.reused, "IMAGE-9");
1170        assert!(!ensured.uploaded);
1171        assert_eq!(ensured.image.name, name);
1172        assert_eq!(ensured.image.version, "1");
1173        assert_eq!(fake.operations(), ["GetMicrovmImage"], "IMAGE-9");
1174        assert!(uploads(&services).is_empty());
1175        assert!(
1176            fake.paths()[0].contains(&paths_encoded(&format!(
1177                "arn:aws:lambda:us-east-1:{ACCOUNT}:microvm-image:{name}"
1178            ))),
1179            "IMAGE-8: the describe names the ARN built from the account: {:?}",
1180            fake.paths()
1181        );
1182    }
1183
1184    fn paths_encoded(arn: &str) -> String {
1185        crate::control::transport::paths::encode_segment(arn)
1186    }
1187
1188    /// **IMAGE-9, wait.** A build already running under the name is waited out and reused;
1189    /// this caller uploads and creates nothing.
1190    #[tokio::test]
1191    async fn a_running_build_is_waited_out_and_reused() {
1192        let name = expected_name();
1193        let fake = Arc::new(FakeControlPlane::new());
1194        fake.answer("GetMicrovmImage", state(&name, "CREATING"))
1195            .answer("GetMicrovmImage", Answer::ok(ready(&name)));
1196        let services = Arc::new(FakeServices::default());
1197        let ensured = sandbox(&fake, &services)
1198            .ensure_image(request())
1199            .await
1200            .expect("waited");
1201        assert!(ensured.reused, "IMAGE-9");
1202        assert_eq!(fake.call_count("CreateMicrovmImage"), 0, "IMAGE-9");
1203        assert!(uploads(&services).is_empty(), "IMAGE-9");
1204    }
1205
1206    /// **IMAGE-8, the build.** No image under the name: the artifact goes to the
1207    /// content-addressed key, the create names that URI and the derived name, and the image
1208    /// this call built is not reused.
1209    #[tokio::test]
1210    async fn an_absent_image_is_uploaded_to_its_key_and_built() {
1211        let name = expected_name();
1212        let fake = Arc::new(FakeControlPlane::new());
1213        fake.answer("GetMicrovmImage", absent())
1214            .answer("GetMicrovmImage", state(&name, "CREATING"))
1215            .answer("GetMicrovmImage", Answer::ok(ready(&name)))
1216            .answer(
1217                "CreateMicrovmImage",
1218                Answer::created(fake::create_image_response(&name)),
1219            );
1220        let services = Arc::new(FakeServices::default());
1221        let mut request = request();
1222        request
1223            .tags
1224            .insert("harbor:task".to_string(), "t".to_string());
1225        let ensured = sandbox(&fake, &services)
1226            .ensure_image(request)
1227            .await
1228            .expect("built");
1229        assert!(!ensured.reused);
1230        assert!(ensured.uploaded);
1231        let key = format!("harbor/images/{name}/artifact.zip");
1232        assert_eq!(
1233            uploads(&services),
1234            [(BUCKET.to_string(), key.clone())],
1235            "IMAGE-8"
1236        );
1237        assert_eq!(ensured.artifact_uri, format!("s3://{BUCKET}/{key}"));
1238        let body = fake.first_body("CreateMicrovmImage");
1239        assert_eq!(body["name"], name.as_str());
1240        assert_eq!(body["codeArtifact"]["uri"], ensured.artifact_uri.as_str());
1241        assert_eq!(body["tags"]["harbor:task"], "t");
1242        assert_eq!(
1243            body["baseImageArn"], "arn:aws:lambda:us-east-1:aws:microvm-image:al2023-1",
1244            "the derived base keeps the managed name"
1245        );
1246        assert_eq!(services.account_calls.load(Ordering::SeqCst), 1);
1247
1248        // The uploaded bytes are the artifact the request describes.
1249        let (_, _, bytes) = services.puts.lock().expect("not poisoned")[0].clone();
1250        let archive = zip::ZipArchive::new(std::io::Cursor::new(bytes)).expect("a zip");
1251        assert_eq!(archive.len(), 2);
1252    }
1253
1254    /// **IMAGE-10, a failure.** A failed image is deleted, the name awaited free, and then
1255    /// rebuilt; the delete comes before the create.
1256    #[tokio::test]
1257    async fn a_failed_image_is_deleted_then_rebuilt() {
1258        let name = expected_name();
1259        let fake = Arc::new(FakeControlPlane::new());
1260        fake.answer("GetMicrovmImage", state(&name, "CREATE_FAILED"))
1261            .answer("GetMicrovmImage", absent())
1262            .answer("GetMicrovmImage", state(&name, "CREATING"))
1263            .answer("GetMicrovmImage", Answer::ok(ready(&name)))
1264            .answer(
1265                "ListMicrovmImageVersions",
1266                Answer::ok(fake::list_versions_response("1")),
1267            )
1268            .answer(
1269                "DeleteMicrovmImage",
1270                Answer::ok(fake::delete_image_response()),
1271            )
1272            .answer(
1273                "CreateMicrovmImage",
1274                Answer::created(fake::create_image_response(&name)),
1275            );
1276        let services = Arc::new(FakeServices::default());
1277        let ensured = sandbox(&fake, &services)
1278            .ensure_image(request())
1279            .await
1280            .expect("rebuilt");
1281        assert!(!ensured.reused);
1282        let operations = fake.operations();
1283        let deleted = operations
1284            .iter()
1285            .position(|op| *op == "DeleteMicrovmImage")
1286            .expect("IMAGE-10: deleted");
1287        let created = operations
1288            .iter()
1289            .position(|op| *op == "CreateMicrovmImage")
1290            .expect("rebuilt");
1291        assert!(deleted < created, "IMAGE-10: {operations:?}");
1292    }
1293
1294    /// **IMAGE-10, force.** A ready image is rebuilt only when the caller forces it.
1295    #[tokio::test]
1296    async fn a_ready_image_is_rebuilt_only_under_force() {
1297        let name = expected_name();
1298        let fake = Arc::new(FakeControlPlane::new());
1299        fake.answer("GetMicrovmImage", Answer::ok(ready(&name)))
1300            .answer("GetMicrovmImage", absent())
1301            .answer("GetMicrovmImage", state(&name, "CREATING"))
1302            .answer("GetMicrovmImage", Answer::ok(ready(&name)))
1303            .answer(
1304                "ListMicrovmImageVersions",
1305                Answer::ok(fake::list_versions_response("1")),
1306            )
1307            .answer(
1308                "DeleteMicrovmImage",
1309                Answer::ok(fake::delete_image_response()),
1310            )
1311            .answer(
1312                "CreateMicrovmImage",
1313                Answer::created(fake::create_image_response(&name)),
1314            );
1315        let services = Arc::new(FakeServices::default());
1316        let mut forced = request();
1317        forced.force = true;
1318        let ensured = sandbox(&fake, &services)
1319            .ensure_image(forced)
1320            .await
1321            .expect("rebuilt");
1322        assert!(!ensured.reused, "IMAGE-10");
1323        assert_eq!(fake.call_count("DeleteMicrovmImage"), 1, "IMAGE-10");
1324        assert_eq!(fake.call_count("CreateMicrovmImage"), 1);
1325    }
1326
1327    /// **IMAGE-10, a deletion under way.** The name is awaited free, then built.
1328    #[tokio::test]
1329    async fn a_deletion_under_way_is_awaited_then_built() {
1330        let name = expected_name();
1331        let fake = Arc::new(FakeControlPlane::new());
1332        fake.answer("GetMicrovmImage", state(&name, "DELETING"))
1333            .answer("GetMicrovmImage", absent())
1334            .answer("GetMicrovmImage", state(&name, "CREATING"))
1335            .answer("GetMicrovmImage", Answer::ok(ready(&name)))
1336            .answer(
1337                "CreateMicrovmImage",
1338                Answer::created(fake::create_image_response(&name)),
1339            );
1340        let services = Arc::new(FakeServices::default());
1341        let ensured = sandbox(&fake, &services)
1342            .ensure_image(request())
1343            .await
1344            .expect("built");
1345        assert!(!ensured.reused);
1346        assert_eq!(
1347            fake.call_count("DeleteMicrovmImage"),
1348            0,
1349            "not this caller's delete"
1350        );
1351    }
1352
1353    /// **IMAGE-11, the race.** The create is refused because a sibling created the name
1354    /// first: describe again, wait for the sibling's build, and return it as reused.
1355    #[tokio::test]
1356    async fn a_refused_create_joins_the_winners_build() {
1357        let name = expected_name();
1358        let fake = Arc::new(FakeControlPlane::new());
1359        fake.answer("GetMicrovmImage", absent())
1360            .answer("GetMicrovmImage", state(&name, "CREATING"))
1361            .answer("GetMicrovmImage", state(&name, "CREATING"))
1362            .answer("GetMicrovmImage", Answer::ok(ready(&name)))
1363            .answer(
1364                "CreateMicrovmImage",
1365                Answer::failure(409, "An image with this name already exists"),
1366            );
1367        let services = Arc::new(FakeServices::default());
1368        let ensured = sandbox(&fake, &services)
1369            .ensure_image(request())
1370            .await
1371            .expect("IMAGE-11: joined");
1372        assert!(
1373            ensured.reused,
1374            "IMAGE-11: this caller's create did not build it"
1375        );
1376        assert_eq!(
1377            ensured.image.state, "CREATED",
1378            "IMAGE-11: returned only once ready"
1379        );
1380        assert_eq!(fake.call_count("CreateMicrovmImage"), 1);
1381    }
1382
1383    /// **IMAGE-11, not a race.** A refused create with nothing under the name afterwards is
1384    /// the create's own failure, surfaced as it came.
1385    #[tokio::test]
1386    async fn a_refused_create_with_no_winner_surfaces_the_refusal() {
1387        let fake = Arc::new(FakeControlPlane::new());
1388        fake.answer("GetMicrovmImage", absent()).answer(
1389            "CreateMicrovmImage",
1390            Answer::failure(400, "buildRoleArn is invalid"),
1391        );
1392        let services = Arc::new(FakeServices::default());
1393        let error = sandbox(&fake, &services)
1394            .ensure_image(request())
1395            .await
1396            .expect_err("no winner");
1397        assert!(
1398            error.to_string().contains("buildRoleArn is invalid"),
1399            "{error}"
1400        );
1401    }
1402
1403    /// **IMAGE-8, once per sandbox.** Two ensures on one sandbox make one account lookup.
1404    #[tokio::test]
1405    async fn the_account_is_resolved_once_per_sandbox() {
1406        let name = expected_name();
1407        let fake = Arc::new(FakeControlPlane::new());
1408        fake.answer("GetMicrovmImage", Answer::ok(ready(&name)));
1409        let services = Arc::new(FakeServices::default());
1410        let mut sandbox = sandbox(&fake, &services);
1411        sandbox.ensure_image(request()).await.expect("first");
1412        sandbox.ensure_image(request()).await.expect("second");
1413        assert_eq!(services.account_calls.load(Ordering::SeqCst), 1, "IMAGE-8");
1414    }
1415
1416    /// The account lookup of the test below: it records whether the other task's timer had
1417    /// fired by the time the ensure asked, then refuses, so the ensure goes no further.
1418    struct TimedLookup {
1419        fired: Arc<std::sync::atomic::AtomicBool>,
1420        seen: Mutex<Option<bool>>,
1421    }
1422
1423    impl BuildServices for TimedLookup {
1424        fn caller_account(&self) -> futures_util::future::BoxFuture<'_, Result<String, Error>> {
1425            *self.seen.lock().expect("not poisoned") = Some(self.fired.load(Ordering::SeqCst));
1426            Box::pin(async { Err(Error::new(ErrorKind::Credentials, "the probe's refusal")) })
1427        }
1428
1429        fn put_object<'a>(
1430            &'a self,
1431            _bucket: &'a str,
1432            _key: &'a str,
1433            _bytes: Vec<u8>,
1434        ) -> futures_util::future::BoxFuture<'a, Result<(), Error>> {
1435            Box::pin(async { panic!("the probe refuses before any upload") })
1436        }
1437    }
1438
1439    /// **`ensure_image` prepares its artifact off the caller's worker (#309).** On a
1440    /// current-thread runtime, a one-millisecond timer another task set before the ensure
1441    /// began fires while a daemon-sized artifact is hashed and zipped, before the ensure asks
1442    /// for its account.
1443    ///
1444    /// Inline, `prepare` holds the runtime's only worker from its first line to that lookup,
1445    /// so no other task runs in between: a concurrent ensure's credential fetch parked
1446    /// behind it is what timed out under instance-role credentials. The daemon here is 2 MiB
1447    /// that deflate can't shrink, like the real one, which takes hundreds of milliseconds to
1448    /// prepare in a debug build, so the timer has that long to fire.
1449    ///
1450    /// **Falsification**: call `prepare` inline in `Sandbox::ensure_image` again and the
1451    /// lookup finds the timer unfired.
1452    #[tokio::test(flavor = "current_thread")]
1453    async fn the_artifact_is_prepared_while_the_callers_other_tasks_run() {
1454        let fired = Arc::new(std::sync::atomic::AtomicBool::new(false));
1455        let lookup = Arc::new(TimedLookup {
1456            fired: Arc::clone(&fired),
1457            seen: Mutex::new(None),
1458        });
1459        let fake = Arc::new(FakeControlPlane::new());
1460        let plane = crate::testing::control_plane(
1461            fake.clone(),
1462            Region::UsEast1,
1463            Arc::new(TestClock::new()),
1464        );
1465        let mut sandbox = Sandbox::with_control_plane(plane).with_build_services(lookup.clone());
1466        let mut request = request();
1467        // xorshift64: bytes with no repeats for deflate to find, from no dependency, behind
1468        // the aarch64 header the preflight asks for.
1469        let mut state = 0x9e37_79b9_7f4a_7c15_u64;
1470        let payload: Vec<u8> = (0..(2 << 20) / 8)
1471            .flat_map(|_| {
1472                state ^= state << 13;
1473                state ^= state >> 7;
1474                state ^= state << 17;
1475                state.to_le_bytes()
1476            })
1477            .collect();
1478        request.binary = aarch64_daemon(&payload);
1479
1480        let timer = tokio::spawn(async move {
1481            tokio::time::sleep(Duration::from_millis(1)).await;
1482            fired.store(true, Ordering::SeqCst);
1483        });
1484        let error = sandbox
1485            .ensure_image(request)
1486            .await
1487            .expect_err("the probe refuses the lookup");
1488        assert!(error.to_string().contains("the probe's refusal"), "{error}");
1489        assert_eq!(
1490            *lookup.seen.lock().expect("not poisoned"),
1491            Some(true),
1492            "the other task's timer fires while the artifact is prepared, not after"
1493        );
1494        assert!(fake.calls().is_empty(), "no control-plane call came first");
1495        timer.await.expect("the timer task");
1496    }
1497
1498    /// Every local refusal costs zero calls — not the account lookup, not a describe, not
1499    /// the upload.
1500    #[tokio::test]
1501    async fn a_locally_refused_request_makes_no_call() {
1502        let fake = Arc::new(FakeControlPlane::new());
1503        let services = Arc::new(FakeServices::default());
1504        let mut sandbox = sandbox(&fake, &services);
1505        for (mutate, cause) in [
1506            (
1507                Box::new(|r: &mut EnsureImageRequest| {
1508                    r.dockerfile = "FROM x\nENTRYPOINT [\"/bin/sh\"]\nCMD [\"/agentd\"]\n".into()
1509                }) as Box<dyn Fn(&mut EnsureImageRequest)>,
1510                "ENTRYPOINT",
1511            ),
1512            (Box::new(|r| r.name_prefix = "//".into()), "prefix"),
1513            (Box::new(|r| r.s3_bucket = String::new()), "bucket"),
1514            (
1515                Box::new(|r| r.build_role_arn = "not-an-arn".into()),
1516                "buildRoleArn",
1517            ),
1518            (
1519                Box::new(|r| r.s3_key_prefix = Some("a\nb".into())),
1520                "key prefix",
1521            ),
1522        ] {
1523            let mut request = request();
1524            mutate(&mut request);
1525            let error = sandbox.ensure_image(request).await.expect_err(cause);
1526            assert_eq!(error.kind(), ErrorKind::InvalidArg, "{error}");
1527            assert!(error.to_string().contains(cause), "{cause}: {error}");
1528        }
1529        assert!(fake.calls().is_empty());
1530        assert_eq!(services.account_calls.load(Ordering::SeqCst), 0);
1531        assert!(uploads(&services).is_empty());
1532    }
1533
1534    /// The build wait honours `wait_timeout`: a build that never settles is a timeout at the
1535    /// caller's deadline, not the 45-minute default.
1536    #[tokio::test]
1537    async fn the_wait_timeout_bounds_the_build_wait() {
1538        let name = expected_name();
1539        let fake = Arc::new(FakeControlPlane::new());
1540        fake.answer("GetMicrovmImage", absent())
1541            .answer("GetMicrovmImage", state(&name, "CREATING"))
1542            .answer(
1543                "CreateMicrovmImage",
1544                Answer::created(fake::create_image_response(&name)),
1545            )
1546            .answer(
1547                "ListMicrovmImageBuilds",
1548                Answer::ok(fake::list_builds_response("IN_PROGRESS")),
1549            )
1550            .answer(
1551                "ListMicrovmImageVersions",
1552                Answer::ok(fake::list_versions_response("1")),
1553            );
1554        let services = Arc::new(FakeServices::default());
1555        let mut request = request();
1556        request.wait_timeout = Some(Duration::from_secs(60));
1557        let error = sandbox(&fake, &services)
1558            .ensure_image(request)
1559            .await
1560            .expect_err("never settles");
1561        assert_eq!(error.kind(), ErrorKind::Timeout, "{error}");
1562        assert!(
1563            fake.call_count("GetMicrovmImage") < 10,
1564            "bounded by 60s, not 45 minutes"
1565        );
1566    }
1567}