Skip to main content

bv_runtime/
runtime.rs

1use std::collections::HashMap;
2use std::fmt;
3use std::path::PathBuf;
4use std::str::FromStr;
5use std::time::Duration;
6
7use bv_core::error::Result;
8use bv_core::manifest::GpuSpec;
9
10// OciRef
11
12#[derive(Debug, Clone)]
13pub struct OciRef {
14    pub registry: String,
15    pub repository: String,
16    pub tag: Option<String>,
17    pub digest: Option<String>,
18}
19
20impl OciRef {
21    pub fn parse(s: &str) -> std::result::Result<Self, String> {
22        s.parse()
23    }
24
25    /// Return the string Docker expects for `docker pull` / `docker run`.
26    /// For docker.io images the registry prefix is stripped so that Docker Hub
27    /// resolves references correctly across all Docker versions.
28    pub fn docker_arg(&self) -> String {
29        if self.registry == "docker.io" {
30            let mut s = self.repository.clone();
31            if let Some(tag) = &self.tag {
32                s.push(':');
33                s.push_str(tag);
34            }
35            if let Some(digest) = &self.digest {
36                s.push('@');
37                s.push_str(digest);
38            }
39            s
40        } else {
41            self.to_string()
42        }
43    }
44}
45
46impl fmt::Display for OciRef {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        write!(f, "{}/{}", self.registry, self.repository)?;
49        if let Some(tag) = &self.tag {
50            write!(f, ":{tag}")?;
51        }
52        if let Some(digest) = &self.digest {
53            write!(f, "@{digest}")?;
54        }
55        Ok(())
56    }
57}
58
59impl FromStr for OciRef {
60    type Err = String;
61
62    fn from_str(raw: &str) -> std::result::Result<Self, Self::Err> {
63        let s = raw
64            .strip_prefix("docker://")
65            .or_else(|| raw.strip_prefix("oci://"))
66            .unwrap_or(raw);
67
68        let (image_part, digest) = if let Some((img, d)) = s.split_once('@') {
69            (img, Some(d.to_string()))
70        } else {
71            (s, None)
72        };
73
74        let (name_part, tag) = if let Some(pos) = image_part.rfind(':') {
75            let before = &image_part[..pos];
76            if before.contains('/') || !before.contains(':') {
77                (&image_part[..pos], Some(image_part[pos + 1..].to_string()))
78            } else {
79                (image_part, None)
80            }
81        } else {
82            (image_part, None)
83        };
84
85        let (registry, repository) = split_registry(name_part);
86
87        Ok(OciRef {
88            registry,
89            repository,
90            tag,
91            digest,
92        })
93    }
94}
95
96fn split_registry(name: &str) -> (String, String) {
97    if let Some(slash_pos) = name.find('/') {
98        let potential_registry = &name[..slash_pos];
99        if potential_registry.contains('.')
100            || potential_registry.contains(':')
101            || potential_registry == "localhost"
102        {
103            return (
104                potential_registry.to_string(),
105                name[slash_pos + 1..].to_string(),
106            );
107        }
108    }
109    ("docker.io".to_string(), name.to_string())
110}
111
112// Supporting types
113
114#[derive(Debug, Clone)]
115pub struct ImageDigest(pub String);
116
117impl fmt::Display for ImageDigest {
118    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
119        f.write_str(&self.0)
120    }
121}
122
123#[derive(Debug, Clone)]
124pub struct ImageMetadata {
125    pub digest: ImageDigest,
126    pub size_bytes: Option<u64>,
127    pub labels: HashMap<String, String>,
128}
129
130#[derive(Debug, Clone)]
131pub struct RuntimeInfo {
132    pub name: String,
133    pub version: String,
134    pub extra: HashMap<String, String>,
135}
136
137#[derive(Debug, Clone)]
138pub struct Mount {
139    pub host_path: PathBuf,
140    pub container_path: PathBuf,
141    pub read_only: bool,
142}
143
144#[derive(Debug, Clone, Default)]
145pub struct GpuProfile {
146    pub spec: Option<GpuSpec>,
147}
148
149#[derive(Debug, Clone)]
150pub struct RunSpec {
151    /// OCI reference of the image to run (may carry a pinned digest).
152    pub image: OciRef,
153    pub command: Vec<String>,
154    pub env: HashMap<String, String>,
155    pub mounts: Vec<Mount>,
156    pub gpu: GpuProfile,
157    pub working_dir: Option<PathBuf>,
158}
159
160#[derive(Debug)]
161pub struct RunOutcome {
162    pub exit_code: i32,
163    pub duration: Duration,
164}
165
166// ProgressReporter
167
168pub trait ProgressReporter: Send + Sync {
169    fn update(&self, message: &str, current: Option<u64>, total: Option<u64>);
170    fn finish(&self, message: &str);
171
172    /// Hide our own spinner/bars while a child process draws to the terminal.
173    /// Implementations should clear their lines until the returned guard is dropped.
174    /// The default is a no-op for reporters that don't draw to a TTY.
175    fn pause(&self) -> Box<dyn PauseGuard + '_> {
176        Box::new(NoopPauseGuard)
177    }
178}
179
180pub trait PauseGuard {}
181
182pub struct NoopPauseGuard;
183impl PauseGuard for NoopPauseGuard {}
184
185pub struct NoopProgress;
186
187impl ProgressReporter for NoopProgress {
188    fn update(&self, _: &str, _: Option<u64>, _: Option<u64>) {}
189    fn finish(&self, _: &str) {}
190}
191
192// ContainerRuntime trait
193
194pub trait ContainerRuntime {
195    fn name(&self) -> &str;
196    fn health_check(&self) -> Result<RuntimeInfo>;
197    fn pull(&self, image: &OciRef, progress: &dyn ProgressReporter) -> Result<ImageDigest>;
198    fn run(&self, spec: &RunSpec) -> Result<RunOutcome>;
199    fn inspect(&self, digest: &ImageDigest) -> Result<ImageMetadata>;
200    /// Check whether `image_ref@digest` is already in the local Docker cache.
201    fn is_locally_available(&self, _image_ref: &str, digest: &str) -> bool {
202        self.inspect(&ImageDigest(digest.to_string())).is_ok()
203    }
204    fn gpu_args(&self, profile: &GpuProfile) -> Vec<String>;
205    fn mount_args(&self, mounts: &[Mount]) -> Vec<String>;
206}
207
208// Tests
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213
214    #[test]
215    fn parse_simple_image() {
216        let r: OciRef = "ubuntu:22.04".parse().unwrap();
217        assert_eq!(r.registry, "docker.io");
218        assert_eq!(r.repository, "ubuntu");
219        assert_eq!(r.tag.as_deref(), Some("22.04"));
220        assert!(r.digest.is_none());
221    }
222
223    #[test]
224    fn parse_with_registry() {
225        let r: OciRef = "ghcr.io/biocontainers/bwa:0.7.17".parse().unwrap();
226        assert_eq!(r.registry, "ghcr.io");
227        assert_eq!(r.repository, "biocontainers/bwa");
228        assert_eq!(r.tag.as_deref(), Some("0.7.17"));
229    }
230
231    #[test]
232    fn parse_with_digest() {
233        let r: OciRef = "ubuntu@sha256:abc123".parse().unwrap();
234        assert_eq!(r.digest.as_deref(), Some("sha256:abc123"));
235        assert!(r.tag.is_none());
236    }
237
238    #[test]
239    fn parse_docker_scheme() {
240        let r: OciRef = "docker://biocontainers/bwa:0.7.17".parse().unwrap();
241        assert_eq!(r.registry, "docker.io");
242        assert_eq!(r.repository, "biocontainers/bwa");
243    }
244
245    #[test]
246    fn docker_arg_strips_docker_io() {
247        let r: OciRef = "ncbi/blast:2.14.0".parse().unwrap();
248        assert_eq!(r.docker_arg(), "ncbi/blast:2.14.0");
249
250        let mut r2: OciRef = "ncbi/blast:2.14.0".parse().unwrap();
251        r2.tag = None;
252        r2.digest = Some("sha256:abc123".into());
253        assert_eq!(r2.docker_arg(), "ncbi/blast@sha256:abc123");
254    }
255
256    #[test]
257    fn docker_arg_keeps_external_registry() {
258        let r: OciRef = "quay.io/biocontainers/blast:2.15.0".parse().unwrap();
259        assert_eq!(r.docker_arg(), "quay.io/biocontainers/blast:2.15.0");
260    }
261}