From d9c804accac48bea222012e62d480d5ba8de9d49 Mon Sep 17 00:00:00 2001
From: Ryan Yin <xiaoyin_c@qq.com>
Date: Sat, 30 Mar 2024 22:42:21 +0800
Subject: [PATCH] feat(host/k8s,lib): nodeLabels & nodeTaints for k8s

---
 hosts/k8s/kubevirt-shoryu/default.nix  |  3 +++
 hosts/k8s/kubevirt-shushou/default.nix |  3 +++
 hosts/k8s/kubevirt-youko/default.nix   |  3 +++
 lib/genK3sAgentModule.nix              | 10 ++++++----
 lib/genK3sServerModule.nix             |  7 ++++---
 5 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/hosts/k8s/kubevirt-shoryu/default.nix b/hosts/k8s/kubevirt-shoryu/default.nix
index 0e5f35dc..11f88f7b 100644
--- a/hosts/k8s/kubevirt-shoryu/default.nix
+++ b/hosts/k8s/kubevirt-shoryu/default.nix
@@ -22,6 +22,9 @@
     # use my own domain & kube-vip's virtual IP for the API server
     # so that the API server can always be accessed even if some nodes are down
     masterHost = "kubevirt-cluster-1.writefor.fun";
+    nodeLabels = [
+      "node-purpose=kubevirt"
+    ];
   };
 in {
   imports =
diff --git a/hosts/k8s/kubevirt-shushou/default.nix b/hosts/k8s/kubevirt-shushou/default.nix
index 7f014315..fe6e782b 100644
--- a/hosts/k8s/kubevirt-shushou/default.nix
+++ b/hosts/k8s/kubevirt-shushou/default.nix
@@ -20,6 +20,9 @@
     # use my own domain & kube-vip's virtual IP for the API server
     # so that the API server can always be accessed even if some nodes are down
     masterHost = "kubevirt-cluster-1.writefor.fun";
+    nodeLabels = [
+      "node-purpose=kubevirt"
+    ];
   };
 in {
   imports =
diff --git a/hosts/k8s/kubevirt-youko/default.nix b/hosts/k8s/kubevirt-youko/default.nix
index 4d818781..6b181caa 100644
--- a/hosts/k8s/kubevirt-youko/default.nix
+++ b/hosts/k8s/kubevirt-youko/default.nix
@@ -20,6 +20,9 @@
     # use my own domain & kube-vip's virtual IP for the API server
     # so that the API server can always be accessed even if some nodes are down
     masterHost = "kubevirt-cluster-1.writefor.fun";
+    nodeLabels = [
+      "node-purpose=kubevirt"
+    ];
   };
 in {
   imports =
diff --git a/lib/genK3sAgentModule.nix b/lib/genK3sAgentModule.nix
index 510dc514..96e627e2 100644
--- a/lib/genK3sAgentModule.nix
+++ b/lib/genK3sAgentModule.nix
@@ -2,6 +2,7 @@
   pkgs,
   masterHost,
   tokenFile,
+  nodeLabels ? [],
   ...
 }: let
   package = pkgs.k3s_1_29;
@@ -15,10 +16,11 @@ in {
     serverAddr = "https://${masterHost}:6443";
     # https://docs.k3s.io/cli/agent
     extraFlags = let
-      flagList = [
-        "--node-label=node-type=worker"
-        "--data-dir /var/lib/rancher/k3s"
-      ];
+      flagList =
+        [
+          "--data-dir /var/lib/rancher/k3s"
+        ]
+        ++ (map (label: "--node-label=${label}") nodeLabels);
     in
       pkgs.lib.concatStringsSep " " flagList;
   };
diff --git a/lib/genK3sServerModule.nix b/lib/genK3sServerModule.nix
index 75817d8c..fc1dece7 100644
--- a/lib/genK3sServerModule.nix
+++ b/lib/genK3sServerModule.nix
@@ -10,7 +10,8 @@
   # this can be a domain name or an IP address(such as kube-vip's virtual IP)
   masterHost,
   clusterInit ? false,
-  addTaints ? false,
+  nodeLabels ? [],
+  nodeTaints ? [],
   ...
 }: let
   package = pkgs.k3s_1_29;
@@ -56,8 +57,8 @@ in {
           "--disable-network-policy"
           "--tls-san=${masterHost}"
         ]
-        # prevent workloads from running on the master
-        ++ (pkgs.lib.optionals addTaints ["--node-taint=CriticalAddonsOnly=true:NoExecute"]);
+        ++ (map (label: "--node-label=${label}") nodeLabels)
+        ++ (map (taint: "--node-taint=${taint}") nodeTaints);
     in
       pkgs.lib.concatStringsSep " " flagList;
   };