[node-labeler] Refactor node labeling to use taints instead of labels

Author: iQQBot

components/node-labeler/cmd/run.go

@@ -50,6 +50,13 @@ const (
 
 	registryFacade = "registry-facade"
 	wsDaemon       = "ws-daemon"
+
+	// Taint keys for different components
+	registryFacadeTaintKey = "gitpod.io/registry-facade-not-ready"
+	wsDaemonTaintKey       = "gitpod.io/ws-daemon-not-ready"
+
+	workspacesRegularLabel  = "gitpod.io/workload_workspace_regular"
+	workspacesHeadlessLabel = "gitpod.io/workload_workspace_headless"
 )
 
 var defaultRequeueTime = time.Second * 10
@@ -61,6 +68,15 @@ var runCmd = &cobra.Command{
 	Run: func(cmd *cobra.Command, args []string) {
 		ctrl.SetLogger(logrusr.New(log.Log))
 
+		kClient, err := client.New(ctrl.GetConfigOrDie(), client.Options{})
+		if err != nil {
+			log.WithError(err).Fatal("unable to create client")
+		}
+
+		if err := initializeLabels(context.Background(), kClient); err != nil {
+			log.WithError(err).Fatal("failed to initialize labels")
+		}
+
 		mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
 			Scheme:                 scheme,
 			HealthProbeBindAddress: ":8086",
@@ -84,11 +100,6 @@ var runCmd = &cobra.Command{
 			log.WithError(err).Fatal("unable to start node-labeler")
 		}
 
-		kClient, err := client.New(ctrl.GetConfigOrDie(), client.Options{})
-		if err != nil {
-			log.WithError(err).Fatal("unable to create client")
-		}
-
 		r := &PodReconciler{
 			kClient,
 		}
@@ -198,21 +209,18 @@ func (r *PodReconciler) Reconcile(ctx context.Context, req reconcile.Request) (r
 	}
 
 	var (
-		ipAddress     string
-		port          string
-		component     string
-		labelToUpdate string
+		ipAddress string
+		port      string
+		taintKey  string
 	)
 
 	switch {
 	case strings.HasPrefix(pod.Name, registryFacade):
-		component = registryFacade
-		labelToUpdate = fmt.Sprintf(registryFacadeLabel, namespace)
+		taintKey = registryFacadeTaintKey
 		ipAddress = pod.Status.HostIP
 		port = strconv.Itoa(registryFacadePort)
 	case strings.HasPrefix(pod.Name, wsDaemon):
-		component = wsDaemon
-		labelToUpdate = fmt.Sprintf(wsdaemonLabel, namespace)
+		taintKey = wsDaemonTaintKey
 		ipAddress = pod.Status.PodIP
 		port = strconv.Itoa(wsdaemonPort)
 	default:
@@ -222,17 +230,17 @@ func (r *PodReconciler) Reconcile(ctx context.Context, req reconcile.Request) (r
 
 	if !pod.ObjectMeta.DeletionTimestamp.IsZero() {
 		// the pod is being removed.
-		// remove the component label from the node
+		// add the taint to the node
 		time.Sleep(1 * time.Second)
-		err := updateLabel(labelToUpdate, false, nodeName, r)
+		err := updateNodeTaint(taintKey, true, nodeName, r)
 		if err != nil {
 			// this is a edge case when cluster-autoscaler removes a node
 			// (all the running pods will be removed after that)
 			if errors.IsNotFound(err) {
 				return reconcile.Result{}, nil
 			}
 
-			log.WithError(err).Error("removing node label")
+			log.WithError(err).Error("adding node taint")
 			return reconcile.Result{RequeueAfter: defaultRequeueTime}, err
 		}
 
@@ -250,8 +258,17 @@ func (r *PodReconciler) Reconcile(ctx context.Context, req reconcile.Request) (r
 		return reconcile.Result{}, fmt.Errorf("obtaining node %s: %w", nodeName, err)
 	}
 
-	if labelValue, exists := node.Labels[labelToUpdate]; exists && labelValue == "true" {
-		// nothing to do, the label already exists.
+	// Check if taint exists
+	taintExists := false
+	for _, taint := range node.Spec.Taints {
+		if taint.Key == taintKey {
+			taintExists = true
+			break
+		}
+	}
+
+	if !taintExists {
+		// nothing to do, the taint doesn't exist.
 		return reconcile.Result{}, nil
 	}
 
@@ -261,7 +278,7 @@ func (r *PodReconciler) Reconcile(ctx context.Context, req reconcile.Request) (r
 		return reconcile.Result{RequeueAfter: defaultRequeueTime}, nil
 	}
 
-	if component == registryFacade {
+	if strings.HasPrefix(pod.Name, registryFacade) {
 		err = checkRegistryFacade(ipAddress, port)
 		if err != nil {
 			log.WithError(err).Error("checking registry-facade")
@@ -271,15 +288,15 @@ func (r *PodReconciler) Reconcile(ctx context.Context, req reconcile.Request) (r
 		time.Sleep(1 * time.Second)
 	}
 
-	err = updateLabel(labelToUpdate, true, nodeName, r)
+	err = updateNodeTaint(taintKey, false, nodeName, r)
 	if err != nil {
-		log.WithError(err).Error("updating node label")
-		return reconcile.Result{}, fmt.Errorf("trying to add the label: %v", err)
+		log.WithError(err).Error("removing node taint")
+		return reconcile.Result{}, fmt.Errorf("trying to remove the taint: %v", err)
 	}
 
 	readyIn := time.Since(pod.Status.StartTime.Time)
-	NodeLabelerTimeHistVec.WithLabelValues(component).Observe(readyIn.Seconds())
-	NodeLabelerCounterVec.WithLabelValues(component).Inc()
+	NodeLabelerTimeHistVec.WithLabelValues(strings.Split(pod.Name, "-")[0]).Observe(readyIn.Seconds())
+	NodeLabelerCounterVec.WithLabelValues(strings.Split(pod.Name, "-")[0]).Inc()
 
 	return reconcile.Result{}, nil
 }
@@ -485,7 +502,7 @@ func (c *NodeScaledownAnnotationController) updateNodeAnnotation(ctx context.Con
 	})
 }
 
-func updateLabel(label string, add bool, nodeName string, client client.Client) error {
+func updateNodeTaint(taintKey string, add bool, nodeName string, client client.Client) error {
 	return retry.RetryOnConflict(retry.DefaultBackoff, func() error {
 		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		defer cancel()
@@ -496,12 +513,36 @@ func updateLabel(label string, add bool, nodeName string, client client.Client)
 			return err
 		}
 
+		// Create or remove taint
 		if add {
-			node.Labels[label] = "true"
-			log.WithField("label", label).WithField("node", nodeName).Info("adding label to node")
+			// Add taint if it doesn't exist
+			taintExists := false
+			for _, taint := range node.Spec.Taints {
+				if taint.Key == taintKey {
+					taintExists = true
+					break
+				}
+			}
+			if !taintExists {
+				node.Spec.Taints = append(node.Spec.Taints, corev1.Taint{
+					Key:    taintKey,
+					Value:  "true",
+					Effect: corev1.TaintEffectNoSchedule,
+				})
+				log.WithField("taint", taintKey).WithField("node", nodeName).Info("adding taint to node")
+			}
 		} else {
-			delete(node.Labels, label)
-			log.WithField("label", label).WithField("node", nodeName).Info("removing label from node")
+			// Remove taint if it exists
+			newTaints := make([]corev1.Taint, 0)
+			for _, taint := range node.Spec.Taints {
+				if taint.Key != taintKey {
+					newTaints = append(newTaints, taint)
+				}
+			}
+			if len(newTaints) != len(node.Spec.Taints) {
+				node.Spec.Taints = newTaints
+				log.WithField("taint", taintKey).WithField("node", nodeName).Info("removing taint from node")
+			}
 		}
 
 		err = client.Update(ctx, &node)
@@ -569,3 +610,72 @@ func newDefaultTransport() *http.Transport {
 		DisableKeepAlives:     true,
 	}
 }
+
+func initializeLabels(ctx context.Context, kClient client.Client) error {
+	log.Info("initializing labels on nodes")
+
+	var nodes corev1.NodeList
+	if err := kClient.List(ctx, &nodes); err != nil {
+		return fmt.Errorf("failed to list nodes: %w", err)
+	}
+
+	for _, node := range nodes.Items {
+		if node.Labels == nil {
+			continue
+		}
+		_, isRegularWorkspaceNode := node.Labels[workspacesRegularLabel]
+		_, isHeadlessWorkspaceNode := node.Labels[workspacesHeadlessLabel]
+
+		if isRegularWorkspaceNode || isHeadlessWorkspaceNode {
+			err := updateNodeLabel(node.Name, kClient)
+			if err != nil {
+				log.WithError(err).WithField("node", node.Name).Error("failed to initialize labels on node")
+			}
+		}
+	}
+
+	log.Info("finished initializing labels on nodes")
+	return nil
+}
+
+func updateNodeLabel(nodeName string, client client.Client) error {
+	return retry.RetryOnConflict(retry.DefaultBackoff, func() error {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		var node corev1.Node
+		err := client.Get(ctx, types.NamespacedName{Name: nodeName}, &node)
+		if err != nil {
+			return err
+		}
+
+		registryFacadeLabelForNamespace := fmt.Sprintf(registryFacadeLabel, namespace)
+		wsDaemonLabelForNamespace := fmt.Sprintf(wsdaemonLabel, namespace)
+
+		needUpdate := false
+
+		if node.Labels == nil {
+			node.Labels = make(map[string]string)
+		}
+
+		if v := node.Labels[registryFacadeLabelForNamespace]; v != "true" {
+			needUpdate = true
+		}
+		if v := node.Labels[wsDaemonLabelForNamespace]; v != "true" {
+			needUpdate = true
+		}
+
+		if !needUpdate {
+			return nil
+		}
+		node.Labels[registryFacadeLabelForNamespace] = "true"
+		node.Labels[wsDaemonLabelForNamespace] = "true"
+
+		err = client.Update(ctx, &node)
+		if err != nil {
+			return err
+		}
+
+		return nil
+	})
+}