From 4aa91bd94770c6dadba04a33116f9995b398a4e9 Mon Sep 17 00:00:00 2001
From: Lucian Petrut <lpetrut@cloudbasesolutions.com>
Date: Wed, 11 Dec 2024 19:01:19 +0200
Subject: [PATCH] Add worker bootstrap-node-taints setting (#215)

* Add worker bootstrap-node-taints setting

The control-plane charm already allows defining node taints. We'll add the same for worker nodes.
---
 charms/worker/charmcraft.yaml              | 13 +++++++++++++
 charms/worker/k8s/src/charm.py             |  3 +++
 charms/worker/k8s/src/config/extra_args.py | 17 ++++++++++++++++-
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/charms/worker/charmcraft.yaml b/charms/worker/charmcraft.yaml
index 218a7a6c..b73516e0 100644
--- a/charms/worker/charmcraft.yaml
+++ b/charms/worker/charmcraft.yaml
@@ -60,6 +60,19 @@ bases:
 
 config:
   options:
+    bootstrap-node-taints:
+      type: string
+      default: ""
+      description: |
+        Space-separated list of taints to apply to this node at registration time.
+
+        This config is only used at bootstrap time when Kubelet first registers the
+        node with Kubernetes. To change node taints after deploy time, use kubectl
+        instead.
+
+        For more information, see the upstream Kubernetes documentation about
+        taints:
+        https://kubernetes.io/docs/concepts/scheduling-eviction/taint-and-toleration/
     node-labels:
       default: ""
       type: string
diff --git a/charms/worker/k8s/src/charm.py b/charms/worker/k8s/src/charm.py
index fc438a8b..9f686f6d 100755
--- a/charms/worker/k8s/src/charm.py
+++ b/charms/worker/k8s/src/charm.py
@@ -815,6 +815,9 @@ def _join_with_token(self, relation: ops.Relation, token: str, cluster_name: str
             request.config = NodeJoinConfig()
             config.extra_args.craft(self.config, request.config, cluster_name)
 
+            bootstrap_node_taints = str(self.config["bootstrap-node-taints"] or "").strip().split()
+            config.extra_args.taint_worker(request.config, bootstrap_node_taints)
+
         self.api_manager.join_cluster(request)
         log.info("Joined %s(%s)", self.unit, node_name)
 
diff --git a/charms/worker/k8s/src/config/extra_args.py b/charms/worker/k8s/src/config/extra_args.py
index b3be1d8f..0166a05e 100644
--- a/charms/worker/k8s/src/config/extra_args.py
+++ b/charms/worker/k8s/src/config/extra_args.py
@@ -4,7 +4,7 @@
 # Learn more at: https://juju.is/docs/sdk
 
 """Parse extra arguments for Kubernetes components."""
-from typing import Dict, Union
+from typing import Dict, List, Union
 
 import ops
 from charms.k8s.v0.k8sd_api_manager import (
@@ -70,3 +70,18 @@ def craft(
 
     cmd = _parse(src["kubelet-extra-args"])
     dest.extra_node_kubelet_args = cmd
+
+
+def taint_worker(dest: NodeJoinConfig, taints: List[str]):
+    """Apply the specified list of taints to the node join configuration.
+
+    Updates the following attributes of the `config` object:
+        - extra_node_kubelet_args: arguments for kubelet.
+
+    Args:
+        dest (NodeJoinConfig):
+            The configuration object to be updated with extra arguments.
+        taints (List[str]):
+            The list of taints to apply.
+    """
+    dest.extra_node_kubelet_args["--register-with-taints"] = ",".join(taints)