Public release release-e570b80c3a0f

Source commit: e570b80c3a0f4b152a9aa0e7bb8635b3eebcec27

Public tree identity: sha256:cff3388143eb2e9476804cfa1c1246e3a46d485d8aef90f2cdd4df45f2b69abf

crates/clusterflux-core/src/scheduler.rs

@@ -0,0 +1,472 @@
+use std::collections::{BTreeMap, BTreeSet};
+
+use serde::{Deserialize, Serialize};
+use thiserror::Error;
+
+use crate::{
+    ArtifactId, Capability, Digest, EnvironmentRequirements, NodeCapabilities, NodeId, ProjectId,
+    TenantId,
+};
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct NodeDescriptor {
+    pub id: NodeId,
+    pub tenant: TenantId,
+    pub project: ProjectId,
+    pub capabilities: NodeCapabilities,
+    pub cached_environments: BTreeSet<Digest>,
+    pub dependency_caches: BTreeSet<Digest>,
+    pub source_snapshots: BTreeSet<Digest>,
+    pub artifact_locations: BTreeSet<ArtifactId>,
+    pub direct_connectivity: bool,
+    pub online: bool,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct PlacementRequest {
+    pub tenant: TenantId,
+    pub project: ProjectId,
+    pub environment: Option<EnvironmentRequirements>,
+    pub environment_digest: Option<Digest>,
+    #[serde(default)]
+    pub environment_cache_required: bool,
+    pub required_capabilities: BTreeSet<Capability>,
+    pub dependency_cache: Option<Digest>,
+    pub source_snapshot: Option<Digest>,
+    pub required_artifacts: BTreeSet<ArtifactId>,
+    pub quota_available: bool,
+    pub policy_allowed: bool,
+    pub prefer_node: Option<NodeId>,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct Placement {
+    pub node: NodeId,
+    pub score: i64,
+    pub reasons: Vec<String>,
+}
+
+#[derive(Clone, Debug, Error, PartialEq, Eq)]
+#[error("no capable node for placement: {message}")]
+pub struct PlacementError {
+    pub message: String,
+}
+
+pub trait Scheduler {
+    fn place(
+        &self,
+        nodes: &[NodeDescriptor],
+        request: &PlacementRequest,
+    ) -> Result<Placement, PlacementError>;
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct DefaultScheduler;
+
+impl Scheduler for DefaultScheduler {
+    fn place(
+        &self,
+        nodes: &[NodeDescriptor],
+        request: &PlacementRequest,
+    ) -> Result<Placement, PlacementError> {
+        let mut scored = Vec::new();
+        let mut rejection_counts = BTreeMap::<String, usize>::new();
+
+        for node in nodes {
+            match compatibility(node, request) {
+                Ok(mut placement) => {
+                    locality_score(node, request, &mut placement);
+                    scored.push(placement);
+                }
+                Err(reasons) => {
+                    for reason in reasons {
+                        *rejection_counts.entry(reason).or_default() += 1;
+                    }
+                }
+            }
+        }
+
+        scored
+            .into_iter()
+            .max_by_key(|placement| placement.score)
+            .ok_or_else(|| PlacementError {
+                message: rejection_counts
+                    .into_iter()
+                    .map(|(reason, count)| format!("{reason} ({count} node(s))"))
+                    .collect::<Vec<_>>()
+                    .join("; "),
+            })
+    }
+}
+
+fn compatibility(
+    node: &NodeDescriptor,
+    request: &PlacementRequest,
+) -> Result<Placement, Vec<String>> {
+    let mut reasons = Vec::new();
+    if !node.online {
+        reasons.push("node offline".to_owned());
+    }
+    if node.tenant != request.tenant {
+        reasons.push("tenant mismatch".to_owned());
+    }
+    if node.project != request.project {
+        reasons.push("project mismatch".to_owned());
+    }
+    if !request.quota_available {
+        reasons.push("quota unavailable for placement".to_owned());
+    }
+    if !request.policy_allowed {
+        reasons.push("policy denied placement".to_owned());
+    }
+    for capability in &request.required_capabilities {
+        if !node.capabilities.capabilities.contains(capability) {
+            reasons.push(format!("missing capability {capability:?}"));
+        }
+    }
+    if let Some(environment) = &request.environment {
+        if let Some(required_os) = &environment.os {
+            if &node.capabilities.os != required_os {
+                reasons.push(format!("environment requires os {required_os:?}"));
+            }
+        }
+        if let Some(required_arch) = &environment.arch {
+            if &node.capabilities.arch != required_arch {
+                reasons.push(format!("environment requires arch {required_arch}"));
+            }
+        }
+        for capability in &environment.capabilities {
+            if !node.capabilities.capabilities.contains(capability) {
+                reasons.push(format!("environment requires capability {capability:?}"));
+            }
+        }
+    }
+    if request.environment_cache_required {
+        match request.environment_digest.as_ref() {
+            Some(digest) if !node.cached_environments.contains(digest) => {
+                reasons.push(format!(
+                    "required named environment cache {digest} is unavailable"
+                ));
+            }
+            None => reasons.push("required named environment cache digest is missing".to_owned()),
+            Some(_) => {}
+        }
+    }
+    let source_transfer_required = request
+        .source_snapshot
+        .as_ref()
+        .is_some_and(|digest| !node.source_snapshots.contains(digest));
+    if source_transfer_required && !node.direct_connectivity {
+        reasons.push("source snapshot unavailable and direct connectivity unavailable".to_owned());
+    }
+    let missing_artifacts = request
+        .required_artifacts
+        .iter()
+        .filter(|artifact| !node.artifact_locations.contains(*artifact))
+        .count();
+    if missing_artifacts > 0 && !node.direct_connectivity {
+        reasons.push(format!(
+            "{missing_artifacts} required artifact(s) unavailable and direct connectivity unavailable"
+        ));
+    }
+
+    if reasons.is_empty() {
+        Ok(Placement {
+            node: node.id.clone(),
+            score: 0,
+            reasons: Vec::new(),
+        })
+    } else {
+        Err(reasons)
+    }
+}
+
+fn locality_score(node: &NodeDescriptor, request: &PlacementRequest, placement: &mut Placement) {
+    if request.prefer_node.as_ref() == Some(&node.id) {
+        placement.score += 100;
+        placement.reasons.push("preferred node".to_owned());
+    }
+    if request
+        .environment_digest
+        .as_ref()
+        .is_some_and(|digest| node.cached_environments.contains(digest))
+    {
+        placement.score += 50;
+        placement.reasons.push("warm environment cache".to_owned());
+    }
+    if request
+        .source_snapshot
+        .as_ref()
+        .is_some_and(|digest| node.source_snapshots.contains(digest))
+    {
+        placement.score += 40;
+        placement
+            .reasons
+            .push("source snapshot already local".to_owned());
+    }
+    if request
+        .dependency_cache
+        .as_ref()
+        .is_some_and(|digest| node.dependency_caches.contains(digest))
+    {
+        placement.score += 30;
+        placement.reasons.push("warm dependency cache".to_owned());
+    }
+    let artifact_hits = request
+        .required_artifacts
+        .iter()
+        .filter(|artifact| node.artifact_locations.contains(*artifact))
+        .count() as i64;
+    if artifact_hits > 0 {
+        placement.score += 10 * artifact_hits;
+        placement.reasons.push(format!(
+            "{artifact_hits} required artifact(s) already local"
+        ));
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{EnvironmentBackend, Os};
+
+    use super::*;
+
+    fn node(id: &str, cached_source: bool) -> NodeDescriptor {
+        let source = Digest::sha256("source");
+        NodeDescriptor {
+            id: NodeId::from(id),
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            capabilities: NodeCapabilities {
+                os: Os::Linux,
+                arch: "x86_64".to_owned(),
+                capabilities: BTreeSet::from([
+                    Capability::Command,
+                    Capability::Containers,
+                    Capability::RootlessPodman,
+                ]),
+                environment_backends: BTreeSet::from([EnvironmentBackend::Container]),
+                source_providers: BTreeSet::from(["filesystem".to_owned()]),
+            },
+            cached_environments: BTreeSet::from([Digest::sha256("env")]),
+            dependency_caches: if cached_source {
+                BTreeSet::from([Digest::sha256("deps")])
+            } else {
+                BTreeSet::new()
+            },
+            source_snapshots: if cached_source {
+                BTreeSet::from([source])
+            } else {
+                BTreeSet::new()
+            },
+            artifact_locations: BTreeSet::new(),
+            direct_connectivity: true,
+            online: true,
+        }
+    }
+
+    #[test]
+    fn scheduler_prefers_warm_source_and_environment() {
+        let request = PlacementRequest {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            environment: Some(EnvironmentRequirements::linux_container()),
+            environment_digest: Some(Digest::sha256("env")),
+            environment_cache_required: false,
+            required_capabilities: BTreeSet::from([Capability::Command]),
+            dependency_cache: Some(Digest::sha256("deps")),
+            source_snapshot: Some(Digest::sha256("source")),
+            required_artifacts: BTreeSet::new(),
+            quota_available: true,
+            policy_allowed: true,
+            prefer_node: None,
+        };
+
+        let placement = DefaultScheduler
+            .place(&[node("cold", false), node("warm", true)], &request)
+            .unwrap();
+
+        assert_eq!(placement.node, NodeId::from("warm"));
+        assert!(placement
+            .reasons
+            .iter()
+            .any(|reason| reason.contains("source")));
+        assert!(placement
+            .reasons
+            .iter()
+            .any(|reason| reason.contains("dependency")));
+    }
+
+    #[test]
+    fn scheduler_requires_requested_named_environment_cache() {
+        let request = PlacementRequest {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            environment: None,
+            environment_digest: Some(Digest::sha256("missing-environment")),
+            environment_cache_required: true,
+            required_capabilities: BTreeSet::new(),
+            dependency_cache: None,
+            source_snapshot: None,
+            required_artifacts: BTreeSet::new(),
+            quota_available: true,
+            policy_allowed: true,
+            prefer_node: None,
+        };
+
+        let error = DefaultScheduler
+            .place(&[node("uncached", false)], &request)
+            .unwrap_err();
+
+        assert!(error.message.contains("named environment cache"));
+    }
+
+    #[test]
+    fn scheduler_failure_names_missing_constraint() {
+        let mut request = PlacementRequest {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            environment: None,
+            environment_digest: None,
+            environment_cache_required: false,
+            required_capabilities: BTreeSet::from([Capability::WindowsCommandDev]),
+            dependency_cache: None,
+            source_snapshot: None,
+            required_artifacts: BTreeSet::new(),
+            quota_available: true,
+            policy_allowed: true,
+            prefer_node: None,
+        };
+        request.required_capabilities.insert(Capability::Command);
+
+        let error = DefaultScheduler
+            .place(&[node("linux", false)], &request)
+            .unwrap_err();
+
+        assert!(error.message.contains("WindowsCommandDev"));
+    }
+
+    #[test]
+    fn scheduler_failure_names_environment_constraint() {
+        let request = PlacementRequest {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            environment: Some(EnvironmentRequirements::windows_command_dev()),
+            environment_digest: None,
+            environment_cache_required: false,
+            required_capabilities: BTreeSet::new(),
+            dependency_cache: None,
+            source_snapshot: None,
+            required_artifacts: BTreeSet::new(),
+            quota_available: true,
+            policy_allowed: true,
+            prefer_node: None,
+        };
+
+        let error = DefaultScheduler
+            .place(&[node("linux", false)], &request)
+            .unwrap_err();
+
+        assert!(error.message.contains("environment requires os Windows"));
+        assert!(error
+            .message
+            .contains("environment requires capability WindowsCommandDev"));
+    }
+
+    #[test]
+    fn scheduler_requires_direct_connectivity_when_transfer_is_needed() {
+        let mut disconnected = node("disconnected", false);
+        disconnected.direct_connectivity = false;
+        let mut local = node("local", true);
+        local.direct_connectivity = false;
+        let request = PlacementRequest {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            environment: None,
+            environment_digest: None,
+            environment_cache_required: false,
+            required_capabilities: BTreeSet::from([Capability::Command]),
+            dependency_cache: None,
+            source_snapshot: Some(Digest::sha256("source")),
+            required_artifacts: BTreeSet::new(),
+            quota_available: true,
+            policy_allowed: true,
+            prefer_node: None,
+        };
+
+        let placement = DefaultScheduler
+            .place(&[disconnected, local], &request)
+            .unwrap();
+        assert_eq!(placement.node, NodeId::from("local"));
+
+        let mut disconnected = node("disconnected", false);
+        disconnected.direct_connectivity = false;
+        let error = DefaultScheduler
+            .place(&[disconnected], &request)
+            .unwrap_err();
+
+        assert!(error
+            .message
+            .contains("source snapshot unavailable and direct connectivity unavailable"));
+    }
+
+    #[test]
+    fn scheduler_failure_names_required_artifact_transfer_constraint() {
+        let mut disconnected = node("disconnected", true);
+        disconnected.direct_connectivity = false;
+        let request = PlacementRequest {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            environment: None,
+            environment_digest: None,
+            environment_cache_required: false,
+            required_capabilities: BTreeSet::from([Capability::Command]),
+            dependency_cache: None,
+            source_snapshot: None,
+            required_artifacts: BTreeSet::from([ArtifactId::from("cache")]),
+            quota_available: true,
+            policy_allowed: true,
+            prefer_node: None,
+        };
+
+        let error = DefaultScheduler
+            .place(&[disconnected], &request)
+            .unwrap_err();
+
+        assert!(error
+            .message
+            .contains("1 required artifact(s) unavailable and direct connectivity unavailable"));
+    }
+
+    #[test]
+    fn scheduler_failure_names_quota_and_policy_constraints() {
+        let mut request = PlacementRequest {
+            tenant: TenantId::from("tenant"),
+            project: ProjectId::from("project"),
+            environment: None,
+            environment_digest: None,
+            environment_cache_required: false,
+            required_capabilities: BTreeSet::from([Capability::Command]),
+            dependency_cache: None,
+            source_snapshot: None,
+            required_artifacts: BTreeSet::new(),
+            quota_available: false,
+            policy_allowed: true,
+            prefer_node: None,
+        };
+
+        let error = DefaultScheduler
+            .place(&[node("linux", false)], &request)
+            .unwrap_err();
+
+        assert!(error.message.contains("quota unavailable for placement"));
+
+        request.quota_available = true;
+        request.policy_allowed = false;
+        let error = DefaultScheduler
+            .place(&[node("linux", false)], &request)
+            .unwrap_err();
+
+        assert!(error.message.contains("policy denied placement"));
+    }
+}