Skip to main content

mvm_core/
template.rs

1use sha2::Digest;
2
3use serde::{Deserialize, Serialize};
4
5/// Current schema version for persisted state files.
6pub const CURRENT_SCHEMA_VERSION: u32 = 1;
7
8fn default_schema_version() -> u32 {
9    CURRENT_SCHEMA_VERSION
10}
11
12/// Complete template configuration that can define multiple variants/roles.
13/// Typically loaded from a TOML file.
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct TemplateConfig {
16    /// Optional base name used when a variant omits `name`.
17    #[serde(default)]
18    pub template_id: String,
19    pub flake_ref: String,
20    /// Default profile if a variant omits it.
21    #[serde(default = "default_profile")]
22    pub profile: String,
23    pub variants: Vec<TemplateVariant>,
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct TemplateVariant {
28    /// Template name for this variant; if empty, falls back to `<template_id>-<role>`.
29    #[serde(default)]
30    pub name: String,
31    pub role: String,
32    #[serde(default = "default_profile")]
33    pub profile: String,
34    pub vcpus: u8,
35    pub mem_mib: u32,
36    #[serde(default)]
37    pub data_disk_mib: u32,
38}
39
40fn default_profile() -> String {
41    "minimal".to_string()
42}
43
44/// Global template definition (tenant-agnostic base image).
45#[derive(Debug, Clone, Serialize, Deserialize)]
46pub struct TemplateSpec {
47    /// Schema version for forward-compatible migrations. Current: 1.
48    #[serde(default = "default_schema_version")]
49    pub schema_version: u32,
50    pub template_id: String,
51    pub flake_ref: String,
52    pub profile: String,
53    pub role: String,
54    pub vcpus: u8,
55    pub mem_mib: u32,
56    pub data_disk_mib: u32,
57    pub created_at: String,
58    pub updated_at: String,
59}
60
61/// Path helpers
62pub fn templates_base_dir() -> String {
63    format!("{}/templates", crate::config::mvm_data_dir())
64}
65
66pub fn template_dir(template_id: &str) -> String {
67    format!("{}/{}", templates_base_dir(), template_id)
68}
69
70pub fn template_spec_path(template_id: &str) -> String {
71    format!("{}/template.json", template_dir(template_id))
72}
73
74/// Artifacts base dir for a template.
75pub fn template_artifacts_dir(template_id: &str) -> String {
76    format!("{}/artifacts", template_dir(template_id))
77}
78
79/// Specific revision dir for a template.
80pub fn template_revision_dir(template_id: &str, revision: &str) -> String {
81    format!("{}/{}", template_artifacts_dir(template_id), revision)
82}
83
84/// Symlink to current revision.
85pub fn template_current_symlink(template_id: &str) -> String {
86    format!("{}/current", template_dir(template_id))
87}
88
89/// Snapshot directory within a template revision.
90pub fn template_snapshot_dir(template_id: &str, revision: &str) -> String {
91    format!("{}/snapshot", template_revision_dir(template_id, revision))
92}
93
94/// Metadata about a template's pre-built Firecracker snapshot.
95///
96/// Created by `template build --snapshot` after booting the VM and
97/// waiting for the service to become healthy. Used by `run --template`
98/// to restore the VM instantly instead of cold-booting.
99#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
100pub struct SnapshotInfo {
101    pub created_at: String,
102    pub vmstate_size_bytes: u64,
103    pub mem_size_bytes: u64,
104    /// Boot args used when the snapshot was created (must match on restore).
105    pub boot_args: String,
106    /// vCPU count at snapshot time (must match on restore).
107    pub vcpus: u8,
108    /// Memory MiB at snapshot time (must match on restore).
109    pub mem_mib: u32,
110}
111
112/// Describes what kind of pre-built artifact a template provides.
113///
114/// All backends support `Image` (cold-boot from rootfs). Only backends
115/// with `capabilities().snapshots == true` (e.g. Firecracker) support
116/// `Snapshot` (warm-start from memory image).
117#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
118pub enum TemplateKind {
119    /// Pre-built rootfs image only — cold-boot on every start.
120    /// Supported by all backends.
121    Image,
122    /// Pre-built rootfs + Firecracker memory snapshot — warm-start.
123    /// Only supported by backends with snapshot capability.
124    Snapshot(SnapshotInfo),
125}
126
127#[derive(Debug, Clone, Serialize, Deserialize)]
128pub struct TemplateRevision {
129    #[serde(default = "default_schema_version")]
130    pub schema_version: u32,
131    pub revision_hash: String,
132    pub flake_ref: String,
133    pub flake_lock_hash: String,
134    pub artifact_paths: crate::pool::ArtifactPaths,
135    pub built_at: String,
136    pub profile: String,
137    pub role: String,
138    pub vcpus: u8,
139    pub mem_mib: u32,
140    pub data_disk_mib: u32,
141    #[serde(default)]
142    pub snapshot: Option<SnapshotInfo>,
143}
144
145impl TemplateRevision {
146    /// Composite cache key from the three dimensions that define a unique build
147    /// output: flake.lock content, Nix profile, and workload role.
148    pub fn cache_key(&self) -> String {
149        let mut hasher = sha2::Sha256::new();
150        hasher.update(self.flake_lock_hash.as_bytes());
151        hasher.update(b":");
152        hasher.update(self.profile.as_bytes());
153        hasher.update(b":");
154        hasher.update(self.role.as_bytes());
155        format!("{:x}", hasher.finalize())
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162    use crate::pool::ArtifactPaths;
163
164    fn make_revision(flake_lock_hash: &str, profile: &str, role: &str) -> TemplateRevision {
165        TemplateRevision {
166            schema_version: CURRENT_SCHEMA_VERSION,
167            revision_hash: "abc123".to_string(),
168            flake_ref: ".".to_string(),
169            flake_lock_hash: flake_lock_hash.to_string(),
170            artifact_paths: ArtifactPaths {
171                vmlinux: "vmlinux".to_string(),
172                rootfs: "rootfs.ext4".to_string(),
173                fc_base_config: "fc-base.json".to_string(),
174                initrd: None,
175                sizes: None,
176            },
177            built_at: "2025-01-01T00:00:00Z".to_string(),
178            profile: profile.to_string(),
179            role: role.to_string(),
180            vcpus: 2,
181            mem_mib: 1024,
182            data_disk_mib: 0,
183            snapshot: None,
184        }
185    }
186
187    #[test]
188    fn same_inputs_same_cache_key() {
189        let a = make_revision("lock1", "minimal", "worker");
190        let b = make_revision("lock1", "minimal", "worker");
191        assert_eq!(a.cache_key(), b.cache_key());
192    }
193
194    #[test]
195    fn different_profile_different_cache_key() {
196        let a = make_revision("lock1", "minimal", "worker");
197        let b = make_revision("lock1", "full", "worker");
198        assert_ne!(a.cache_key(), b.cache_key());
199    }
200
201    #[test]
202    fn different_role_different_cache_key() {
203        let a = make_revision("lock1", "minimal", "worker");
204        let b = make_revision("lock1", "minimal", "gateway");
205        assert_ne!(a.cache_key(), b.cache_key());
206    }
207
208    #[test]
209    fn different_flake_different_cache_key() {
210        let a = make_revision("lock1", "minimal", "worker");
211        let b = make_revision("lock2", "minimal", "worker");
212        assert_ne!(a.cache_key(), b.cache_key());
213    }
214
215    #[test]
216    fn cache_key_depends_on_flake_lock_not_revision_hash() {
217        let mut a = make_revision("same-lock", "minimal", "worker");
218        a.revision_hash = "rev-aaa".to_string();
219        let mut b = make_revision("same-lock", "minimal", "worker");
220        b.revision_hash = "rev-zzz".to_string();
221        // Different revision hashes but same flake_lock/profile/role → same cache key
222        assert_eq!(a.cache_key(), b.cache_key());
223    }
224
225    #[test]
226    fn snapshot_info_serde_roundtrip() {
227        let info = SnapshotInfo {
228            created_at: "2025-03-01T00:00:00Z".to_string(),
229            vmstate_size_bytes: 1024,
230            mem_size_bytes: 1048576,
231            boot_args: "root=/dev/vda rw init=/init console=ttyS0".to_string(),
232            vcpus: 2,
233            mem_mib: 1024,
234        };
235        let json = serde_json::to_string(&info).unwrap();
236        let back: SnapshotInfo = serde_json::from_str(&json).unwrap();
237        assert_eq!(back.vcpus, 2);
238        assert_eq!(back.mem_mib, 1024);
239        assert_eq!(back.vmstate_size_bytes, 1024);
240    }
241
242    #[test]
243    fn revision_without_snapshot_deserializes() {
244        let json = r#"{
245            "revision_hash": "abc",
246            "flake_ref": ".",
247            "flake_lock_hash": "lock1",
248            "artifact_paths": {
249                "vmlinux": "vmlinux",
250                "rootfs": "rootfs.ext4",
251                "fc_base_config": "fc-base.json"
252            },
253            "built_at": "2025-01-01T00:00:00Z",
254            "profile": "minimal",
255            "role": "worker",
256            "vcpus": 2,
257            "mem_mib": 1024,
258            "data_disk_mib": 0
259        }"#;
260        let rev: TemplateRevision = serde_json::from_str(json).unwrap();
261        assert!(rev.snapshot.is_none());
262    }
263
264    #[test]
265    fn revision_with_snapshot_deserializes() {
266        let rev = make_revision("lock1", "minimal", "worker");
267        let mut rev = rev;
268        rev.snapshot = Some(SnapshotInfo {
269            created_at: "2025-03-01T00:00:00Z".to_string(),
270            vmstate_size_bytes: 512,
271            mem_size_bytes: 2048,
272            boot_args: "console=ttyS0".to_string(),
273            vcpus: 2,
274            mem_mib: 1024,
275        });
276        let json = serde_json::to_string(&rev).unwrap();
277        let back: TemplateRevision = serde_json::from_str(&json).unwrap();
278        assert!(back.snapshot.is_some());
279        assert_eq!(back.snapshot.unwrap().mem_size_bytes, 2048);
280    }
281
282    #[test]
283    fn template_snapshot_dir_format() {
284        let dir = template_snapshot_dir("my-tmpl", "abc123");
285        assert!(dir.ends_with("/templates/my-tmpl/artifacts/abc123/snapshot"));
286    }
287
288    #[test]
289    fn template_kind_image_serde_roundtrip() {
290        let kind = TemplateKind::Image;
291        let json = serde_json::to_string(&kind).unwrap();
292        let parsed: TemplateKind = serde_json::from_str(&json).unwrap();
293        assert_eq!(parsed, TemplateKind::Image);
294    }
295
296    #[test]
297    fn template_kind_snapshot_serde_roundtrip() {
298        let snap = SnapshotInfo {
299            created_at: "2025-03-01T00:00:00Z".to_string(),
300            vmstate_size_bytes: 1024,
301            mem_size_bytes: 2048,
302            boot_args: "console=ttyS0".to_string(),
303            vcpus: 2,
304            mem_mib: 512,
305        };
306        let kind = TemplateKind::Snapshot(snap.clone());
307        let json = serde_json::to_string(&kind).unwrap();
308        let parsed: TemplateKind = serde_json::from_str(&json).unwrap();
309        assert_eq!(parsed, TemplateKind::Snapshot(snap));
310    }
311}