Publish Clusterflux 1d0b6fa filtered source

crates/clusterflux-core/src/vfs.rs

@@ -0,0 +1,241 @@
+use std::collections::BTreeMap;
+
+use serde::{Deserialize, Serialize};
+use thiserror::Error;
+
+use crate::{Digest, NodeId, TaskInstanceId};
+
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+pub struct VfsPath(String);
+
+impl VfsPath {
+    pub fn new(path: impl Into<String>) -> Result<Self, VfsError> {
+        let path = path.into();
+        if !path.starts_with("/vfs/") {
+            return Err(VfsError::InvalidPath(path));
+        }
+        Ok(Self(path))
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct VfsObject {
+    pub path: VfsPath,
+    pub digest: Digest,
+    pub size: u64,
+    pub producer: TaskInstanceId,
+    pub node: NodeId,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct VfsManifest {
+    pub epoch: u64,
+    pub producer: TaskInstanceId,
+    pub node: NodeId,
+    pub objects: BTreeMap<VfsPath, VfsObject>,
+    pub large_bytes_uploaded: bool,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum SyncPolicy {
+    MetadataOnly,
+    ExplicitNode(NodeId),
+    ExplicitStore(String),
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum VfsSyncDecision {
+    NoBytesMoved,
+    MoveBytesToNode(NodeId),
+    MoveBytesToStore(String),
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum ReuseDecision {
+    SameNodeZeroCopy,
+    NeedsTransfer { from: NodeId, to: NodeId },
+    Unavailable,
+}
+
+#[derive(Clone, Debug, Error, PartialEq, Eq)]
+pub enum VfsError {
+    #[error("VFS path must start with /vfs/: {0}")]
+    InvalidPath(String),
+    #[error("path is not visible in the published VFS manifest: {0}")]
+    NotVisible(String),
+}
+
+#[derive(Clone, Debug)]
+pub struct VfsOverlay {
+    task: TaskInstanceId,
+    node: NodeId,
+    epoch: u64,
+    pending: BTreeMap<VfsPath, VfsObject>,
+    published: BTreeMap<VfsPath, VfsObject>,
+}
+
+impl VfsOverlay {
+    pub fn new(task: TaskInstanceId, node: NodeId) -> Self {
+        Self {
+            task,
+            node,
+            epoch: 0,
+            pending: BTreeMap::new(),
+            published: BTreeMap::new(),
+        }
+    }
+
+    pub fn write(&mut self, path: VfsPath, digest: Digest, size: u64) -> VfsObject {
+        let object = VfsObject {
+            path: path.clone(),
+            digest,
+            size,
+            producer: self.task.clone(),
+            node: self.node.clone(),
+        };
+        self.pending.insert(path, object.clone());
+        object
+    }
+
+    pub fn flush(&mut self) -> VfsManifest {
+        self.epoch += 1;
+        self.published.append(&mut self.pending);
+        VfsManifest {
+            epoch: self.epoch,
+            producer: self.task.clone(),
+            node: self.node.clone(),
+            objects: self.published.clone(),
+            large_bytes_uploaded: false,
+        }
+    }
+
+    pub fn sync(&self, policy: SyncPolicy) -> VfsSyncDecision {
+        match policy {
+            SyncPolicy::MetadataOnly => VfsSyncDecision::NoBytesMoved,
+            SyncPolicy::ExplicitNode(node) => VfsSyncDecision::MoveBytesToNode(node),
+            SyncPolicy::ExplicitStore(store) => VfsSyncDecision::MoveBytesToStore(store),
+        }
+    }
+
+    pub fn read_published<'a>(
+        manifest: &'a VfsManifest,
+        path: &VfsPath,
+    ) -> Result<&'a VfsObject, VfsError> {
+        manifest
+            .objects
+            .get(path)
+            .ok_or_else(|| VfsError::NotVisible(path.as_str().to_owned()))
+    }
+
+    pub fn reuse_for_consumer(
+        manifest: &VfsManifest,
+        path: &VfsPath,
+        consumer_node: &NodeId,
+    ) -> ReuseDecision {
+        let Some(object) = manifest.objects.get(path) else {
+            return ReuseDecision::Unavailable;
+        };
+        if &object.node == consumer_node {
+            ReuseDecision::SameNodeZeroCopy
+        } else {
+            ReuseDecision::NeedsTransfer {
+                from: object.node.clone(),
+                to: consumer_node.clone(),
+            }
+        }
+    }
+
+    pub fn discard_unflushed(&mut self) {
+        self.pending.clear();
+    }
+
+    pub fn pending_len(&self) -> usize {
+        self.pending.len()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn path() -> VfsPath {
+        VfsPath::new("/vfs/artifacts/app").unwrap()
+    }
+
+    #[test]
+    fn flush_publishes_manifest_without_large_byte_upload() {
+        let mut overlay = VfsOverlay::new(TaskInstanceId::from("task"), NodeId::from("node-a"));
+        overlay.write(path(), Digest::sha256("binary"), 6);
+
+        let manifest = overlay.flush();
+
+        assert_eq!(manifest.epoch, 1);
+        assert!(!manifest.large_bytes_uploaded);
+        assert!(manifest.objects.contains_key(&path()));
+    }
+
+    #[test]
+    fn downstream_task_can_read_after_flush_but_not_before() {
+        let mut overlay = VfsOverlay::new(TaskInstanceId::from("task"), NodeId::from("node-a"));
+        overlay.write(path(), Digest::sha256("binary"), 6);
+
+        let empty = VfsManifest {
+            epoch: 0,
+            producer: TaskInstanceId::from("task"),
+            node: NodeId::from("node-a"),
+            objects: BTreeMap::new(),
+            large_bytes_uploaded: false,
+        };
+        assert!(VfsOverlay::read_published(&empty, &path()).is_err());
+
+        let manifest = overlay.flush();
+        assert!(VfsOverlay::read_published(&manifest, &path()).is_ok());
+    }
+
+    #[test]
+    fn sync_is_explicit_and_policy_driven() {
+        let overlay = VfsOverlay::new(TaskInstanceId::from("task"), NodeId::from("node-a"));
+
+        assert_eq!(
+            overlay.sync(SyncPolicy::MetadataOnly),
+            VfsSyncDecision::NoBytesMoved
+        );
+        assert_eq!(
+            overlay.sync(SyncPolicy::ExplicitStore("s3://bucket/app".to_owned())),
+            VfsSyncDecision::MoveBytesToStore("s3://bucket/app".to_owned())
+        );
+    }
+
+    #[test]
+    fn same_node_reuse_avoids_transfer() {
+        let mut overlay = VfsOverlay::new(TaskInstanceId::from("task"), NodeId::from("node-a"));
+        overlay.write(path(), Digest::sha256("binary"), 6);
+        let manifest = overlay.flush();
+
+        assert_eq!(
+            VfsOverlay::reuse_for_consumer(&manifest, &path(), &NodeId::from("node-a")),
+            ReuseDecision::SameNodeZeroCopy
+        );
+        assert_eq!(
+            VfsOverlay::reuse_for_consumer(&manifest, &path(), &NodeId::from("node-b")),
+            ReuseDecision::NeedsTransfer {
+                from: NodeId::from("node-a"),
+                to: NodeId::from("node-b")
+            }
+        );
+    }
+
+    #[test]
+    fn unflushed_task_local_changes_can_be_discarded() {
+        let mut overlay = VfsOverlay::new(TaskInstanceId::from("task"), NodeId::from("node-a"));
+        overlay.write(path(), Digest::sha256("binary"), 6);
+
+        overlay.discard_unflushed();
+
+        assert_eq!(overlay.pending_len(), 0);
+    }
+}