Merge pull request 'feat: custom evaluation' (#1) from custom-activation into main

Reviewed-on: DGNum/colmena#1

src/error.rs

@@ -64,6 +64,12 @@ pub enum ColmenaError {
     #[snafu(display("Don't know how to connect to the node"))]
     NoTargetHost,
 
+    #[snafu(display(
+        "Don't know how to deploy node: {} -- does your system type support deployment?",
+        node_name
+    ))]
+    UndeployableHost { node_name: String },
+
     #[snafu(display("Node name cannot be empty"))]
     EmptyNodeName,
 

src/nix/hive/eval.nix

@@ -38,19 +38,25 @@ let
     else if uncheckedHive ? network then uncheckedHive.network
     else {};
 
+  uncheckedRegistries = if uncheckedHive ? registry then uncheckedHive.registry else {};
+
   # The final hive will always have the meta key instead of network.
   hive = let
     userMeta = (lib.modules.evalModules {
       modules = [ colmenaOptions.metaOptions uncheckedUserMeta ];
     }).config;
 
+    registry = (lib.modules.evalModules {
+      modules = [ colmenaOptions.registryOptions  { registry = uncheckedRegistries; } ];
+    }).config.registry;
+
     mergedHive =
       assert lib.assertMsg (!(uncheckedHive ? __schema)) ''
         You cannot pass in an already-evaluated Hive into the evaluator.
 
         Hint: Use the `colmenaHive` output instead of `colmena`.
       '';
-      removeAttrs (defaultHive // uncheckedHive) [ "meta" "network" ];
+      removeAttrs (defaultHive // uncheckedHive) [ "meta" "network" "registry" ];
 
     meta = {
       meta =
@@ -58,7 +64,7 @@ let
         then userMeta // { nixpkgs = <nixpkgs>; }
         else userMeta;
     };
-  in mergedHive // meta;
+  in mergedHive // meta // { inherit registry; };
 
   configsFor = node: let
     nodeConfig = hive.${node};
@@ -112,14 +118,23 @@ let
   in mkNixpkgs "meta.nixpkgs" nixpkgsConf;
 
   lib = nixpkgs.lib;
-  reservedNames = [ "defaults" "network" "meta" ];
+  reservedNames = [ "defaults" "network" "meta" "registry" ];
 
-  evalNode = name: configs: let
+  evalNode = name: configs:
+  # Some help on error messages.
+  assert (lib.assertMsg (lib.hasAttrByPath [ "deployment" "systemType" ] hive.${name})
+  "${name} does not have a deployment system type!");
+  assert (lib.assertMsg (builtins.typeOf hive.registry == "set"))
+    "The hive's registry is not a set, but of type '${builtins.typeOf hive.registry}'";
+  assert (lib.assertMsg (lib.hasAttr hive.${name}.deployment.systemType hive.registry)
+    "${builtins.toJSON (hive.${name}.deployment.systemType)} does not exist in the registry of systems!");
+  let
+    # We cannot use `configs` because we need to access to the raw configuration fragment.
+    inherit (hive.registry.${hive.${name}.deployment.systemType}) evalConfig;
     npkgs =
       if hasAttr name hive.meta.nodeNixpkgs
       then mkNixpkgs "meta.nodeNixpkgs.${name}" hive.meta.nodeNixpkgs.${name}
       else nixpkgs;
-    evalConfig = import (npkgs.path + "/nixos/lib/eval-config.nix");
 
     # Here we need to merge the configurations in meta.nixpkgs
     # and in machine config.
@@ -139,17 +154,19 @@ let
       in
         lib.optional (!hasTypedConfig && length remainingKeys != 0)
         "The following Nixpkgs configuration keys set in meta.nixpkgs will be ignored: ${toString remainingKeys}";
-    };
+      } // lib.optionalAttrs (builtins.hasAttr "localSystem" npkgs || builtins.hasAttr "crossSystem" npkgs) {
+        nixpkgs.localSystem = lib.mkBefore npkgs.localSystem;
+        nixpkgs.crossSystem = lib.mkBefore npkgs.crossSystem;
+      };
   in evalConfig {
-    inherit (npkgs) system;
+    # This doesn't exist for `evalModules` the generic way.
+    # inherit (npkgs) system;
 
     modules = [
       nixpkgsModule
       colmenaModules.assertionModule
-      colmenaModules.keyChownModule
-      colmenaModules.keyServiceModule
       colmenaOptions.deploymentOptions
-      hive.defaults
+      (hive.registry.${hive.${name}.deployment.systemType}.defaults or hive.defaults)
     ] ++ configs;
     specialArgs = {
       inherit name;
@@ -179,6 +196,10 @@ let
     "allowApplyAll"
   ];
 
+  serializableSystemTypeConfigKeys = [
+    "supportsDeployment"
+  ];
+
 in rec {
   # Exported attributes
   __schema = "v0";
@@ -190,5 +211,9 @@ in rec {
   evalSelected =             names: lib.filterAttrs (name: _: elem name names) toplevel;
   evalSelectedDrvPaths =     names: lib.mapAttrs    (_: v: v.drvPath)          (evalSelected names);
   metaConfig = lib.filterAttrs (n: v: elem n metaConfigKeys) hive.meta;
-  introspect = f: f { inherit lib; pkgs = nixpkgs; nodes = uncheckedNodes; };
+  # We cannot perform a `metaConfigKeys`-style simple check here
+  # because registry is arbitrarily deep and may evaluate nixpkgs indirectly.
+  registryConfig = lib.mapAttrs (systemTypeName: systemType:
+    lib.filterAttrs (n: v: elem n serializableSystemTypeConfigKeys) systemType) hive.registry;
+  introspect = f: f { inherit lib; pkgs = nixpkgs; inherit nodes; };
 }

src/nix/hive/mod.rs

@@ -15,7 +15,7 @@ use validator::Validate;
 use super::deployment::TargetNode;
 use super::{
     Flake, MetaConfig, NixExpression, NixFlags, NodeConfig, NodeFilter, NodeName,
-    ProfileDerivation, SerializedNixExpression, StorePath,
+    ProfileDerivation, RegistryConfig, SerializedNixExpression, StorePath,
 };
 use crate::error::{ColmenaError, ColmenaResult};
 use crate::job::JobHandle;
@@ -87,6 +87,8 @@ pub struct Hive {
     nix_options: HashMap<String, String>,
 
     meta_config: OnceCell<MetaConfig>,
+
+    registry_config: OnceCell<RegistryConfig>,
 }
 
 struct NixInstantiate<'hive> {
@@ -140,6 +142,7 @@ impl Hive {
             impure: false,
             nix_options: HashMap::new(),
             meta_config: OnceCell::new(),
+            registry_config: OnceCell::new(),
         })
     }
 
@@ -158,6 +161,17 @@ impl Hive {
             .await
     }
 
+    pub async fn get_registry_config(&self) -> ColmenaResult<&RegistryConfig> {
+        self.registry_config
+            .get_or_try_init(|| async {
+                self.nix_instantiate("hive.registryConfig")
+                    .eval()
+                    .capture_json()
+                    .await
+            })
+            .await
+    }
+
     pub fn set_show_trace(&mut self, value: bool) {
         self.show_trace = value;
     }
@@ -201,6 +215,9 @@ impl Hive {
     ) -> ColmenaResult<HashMap<NodeName, TargetNode>> {
         let mut node_configs = None;
 
+        log::info!("Enumerating systems...");
+        let registry = self.get_registry_config().await?;
+
         log::info!("Enumerating nodes...");
 
         let all_nodes = self.node_names().await?;
@@ -234,6 +251,24 @@ impl Hive {
             self.deployment_info_selected(&selected_nodes).await?
         };
 
+        for node_config in &node_configs {
+            if let Some(system_type) = node_config.1.system_type.as_ref() {
+                let Some(system_config) = registry.systems.get(system_type) else {
+                    // TODO: convert me to proper error?
+                    log::warn!("'{:?}' is not a known system type in the registry, double check your expressions!", system_type);
+                    return Err(ColmenaError::Unknown {
+                        message: "unknown system type".to_string(),
+                    });
+                };
+
+                if !system_config.supports_deployment {
+                    return Err(ColmenaError::UndeployableHost {
+                        node_name: node_config.0.to_string(),
+                    });
+                }
+            }
+        }
+
         let mut targets = HashMap::new();
         let mut n_ssh = 0;
         for node in selected_nodes.into_iter() {

src/nix/hive/options.nix

@@ -98,6 +98,14 @@ with builtins; rec {
   in {
     options = {
       deployment = {
+        systemType = lib.mkOption {
+          description = mdDoc ''
+            System type used for this node, e.g. NixOS.
+          '';
+          default = "nixos";
+          # TODO: enum among all registered systems?
+          type = types.str;
+        };
         targetHost = lib.mkOption {
           description = mdDoc ''
             The target SSH node for deployment.
@@ -218,6 +226,52 @@ with builtins; rec {
       };
     };
   };
+  # Options for a registered system type
+  systemTypeOptions = { name, lib, ... }: let
+    inherit (lib) types;
+    mdDoc = lib.mdDoc or lib.id;
+  in
+  {
+    options = {
+      evalConfig = lib.mkOption {
+        description = mdDoc ''
+          Evaluation function which share the same interface as `nixos/lib/eval-config.nix`
+          which can be tailored to your own usecases or to target another type of system,
+          e.g. nix-darwin.
+        '';
+        type = types.functionTo types.unspecified;
+      };
+      supportsDeployment = lib.mkOption {
+        description = mdDoc ''
+          Whether this system type supports deployment or not.
+
+          If it supports deployment, it needs to have appropriate activation code,
+          refer to how to write custom activators.
+        '';
+        default = name == "nixos";
+        defaultText = "If a NixOS system, then true, otherwise false by default";
+      };
+      defaults = lib.mkOption {
+        description = mdDoc ''
+          Default configuration for that system type.
+        '';
+        type = types.functionTo types.unspecified;
+        default = _: {};
+      };
+    };
+  };
+  registryOptions = { lib, ... }: let
+    inherit (lib) types;
+    mdDoc = lib.mdDoc or lib.id;
+  in
+  {
+    options.registry = lib.mkOption {
+      description = mdDoc ''
+        A registry of all system types.
+      '';
+      type = types.attrsOf (types.submodule systemTypeOptions);
+    };
+  };
   # Hive-wide options
   metaOptions = { lib, ... }: let
     inherit (lib) types;

src/nix/mod.rs

@@ -55,6 +55,9 @@ pub struct NodeName(#[serde(deserialize_with = "NodeName::deserialize")] String)
 
 #[derive(Debug, Clone, Validate, Deserialize)]
 pub struct NodeConfig {
+    #[serde(rename = "systemType")]
+    system_type: Option<String>,
+
     #[serde(rename = "targetHost")]
     target_host: Option<String>,
 
@@ -94,6 +97,18 @@ pub struct MetaConfig {
     pub machines_file: Option<String>,
 }
 
+#[derive(Debug, Clone, Validate, Deserialize)]
+pub struct SystemTypeConfig {
+    #[serde(rename = "supportsDeployment")]
+    pub supports_deployment: bool,
+}
+
+#[derive(Debug, Clone, Validate, Deserialize)]
+pub struct RegistryConfig {
+    #[serde(flatten)]
+    pub systems: HashMap<String, SystemTypeConfig>,
+}
+
 /// Nix CLI flags.
 #[derive(Debug, Clone, Default)]
 pub struct NixFlags {