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
173pub struct NoopProgress;
174
175impl ProgressReporter for NoopProgress {
176    fn update(&self, _: &str, _: Option<u64>, _: Option<u64>) {}
177    fn finish(&self, _: &str) {}
178}
179
180// ContainerRuntime trait
181
182pub trait ContainerRuntime {
183    fn name(&self) -> &str;
184    fn health_check(&self) -> Result<RuntimeInfo>;
185    fn pull(&self, image: &OciRef, progress: &dyn ProgressReporter) -> Result<ImageDigest>;
186    fn run(&self, spec: &RunSpec) -> Result<RunOutcome>;
187    fn inspect(&self, digest: &ImageDigest) -> Result<ImageMetadata>;
188    /// Check whether `image_ref@digest` is already in the local Docker cache.
189    fn is_locally_available(&self, _image_ref: &str, digest: &str) -> bool {
190        self.inspect(&ImageDigest(digest.to_string())).is_ok()
191    }
192    fn gpu_args(&self, profile: &GpuProfile) -> Vec<String>;
193    fn mount_args(&self, mounts: &[Mount]) -> Vec<String>;
194}
195
196// Tests
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201
202    #[test]
203    fn parse_simple_image() {
204        let r: OciRef = "ubuntu:22.04".parse().unwrap();
205        assert_eq!(r.registry, "docker.io");
206        assert_eq!(r.repository, "ubuntu");
207        assert_eq!(r.tag.as_deref(), Some("22.04"));
208        assert!(r.digest.is_none());
209    }
210
211    #[test]
212    fn parse_with_registry() {
213        let r: OciRef = "ghcr.io/biocontainers/bwa:0.7.17".parse().unwrap();
214        assert_eq!(r.registry, "ghcr.io");
215        assert_eq!(r.repository, "biocontainers/bwa");
216        assert_eq!(r.tag.as_deref(), Some("0.7.17"));
217    }
218
219    #[test]
220    fn parse_with_digest() {
221        let r: OciRef = "ubuntu@sha256:abc123".parse().unwrap();
222        assert_eq!(r.digest.as_deref(), Some("sha256:abc123"));
223        assert!(r.tag.is_none());
224    }
225
226    #[test]
227    fn parse_docker_scheme() {
228        let r: OciRef = "docker://biocontainers/bwa:0.7.17".parse().unwrap();
229        assert_eq!(r.registry, "docker.io");
230        assert_eq!(r.repository, "biocontainers/bwa");
231    }
232
233    #[test]
234    fn docker_arg_strips_docker_io() {
235        let r: OciRef = "ncbi/blast:2.14.0".parse().unwrap();
236        assert_eq!(r.docker_arg(), "ncbi/blast:2.14.0");
237
238        let mut r2: OciRef = "ncbi/blast:2.14.0".parse().unwrap();
239        r2.tag = None;
240        r2.digest = Some("sha256:abc123".into());
241        assert_eq!(r2.docker_arg(), "ncbi/blast@sha256:abc123");
242    }
243
244    #[test]
245    fn docker_arg_keeps_external_registry() {
246        let r: OciRef = "quay.io/biocontainers/blast:2.15.0".parse().unwrap();
247        assert_eq!(r.docker_arg(), "quay.io/biocontainers/blast:2.15.0");
248    }
249}