[node-labeler] Refactor node reconciliation and pod health checks

Author: iQQBot

components/node-labeler/cmd/run.go

@@ -26,7 +26,7 @@ import (
 	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
 	_ "k8s.io/client-go/plugin/pkg/client/auth"
 	"k8s.io/client-go/util/retry"
-	"k8s.io/utils/pointer"
+	"k8s.io/utils/ptr"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/builder"
 	"sigs.k8s.io/controller-runtime/pkg/cache"
@@ -35,7 +35,6 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/event"
 	"sigs.k8s.io/controller-runtime/pkg/healthz"
 	"sigs.k8s.io/controller-runtime/pkg/manager"
-	"sigs.k8s.io/controller-runtime/pkg/metrics"
 	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"
 	"sigs.k8s.io/controller-runtime/pkg/predicate"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
@@ -68,15 +67,6 @@ 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",
@@ -88,7 +78,7 @@ var runCmd = &cobra.Command{
 				// default sync period is 10h.
 				// in case node-labeler is restarted and not change happens, we could waste (at least) 20m in a node
 				// that never will run workspaces and the additional nodes cluster-autoscaler adds to compensate
-				SyncPeriod: pointer.Duration(2 * time.Minute),
+				SyncPeriod: ptr.To(time.Duration(2 * time.Minute)),
 			},
 			WebhookServer: webhook.NewServer(webhook.Options{
 				Port: 9443,
@@ -101,7 +91,7 @@ var runCmd = &cobra.Command{
 		}
 
 		r := &PodReconciler{
-			kClient,
+			mgr.GetClient(),
 		}
 
 		componentPredicate, err := predicate.LabelSelectorPredicate(metav1.LabelSelector{
@@ -123,6 +113,25 @@ var runCmd = &cobra.Command{
 		if err != nil {
 			log.WithError(err).Fatal("unable to bind controller watch event handler")
 		}
+		nr := &NodeReconciler{
+			mgr.GetClient(),
+		}
+
+		err = ctrl.NewControllerManagedBy(mgr).
+			Named("node-watcher").
+			For(&corev1.Node{}, builder.WithPredicates(predicate.Or(nr.nodeFilter()))).
+			WithOptions(controller.Options{MaxConcurrentReconciles: 1}).
+			Complete(nr)
+		if err != nil {
+			log.WithError(err).Fatal("unable to bind controller watch event handler")
+		}
+
+		go func() {
+			<-mgr.Elected()
+			if err := nr.reconcileAll(context.Background()); err != nil {
+				log.WithError(err).Fatal("failed to reconcile all nodes")
+			}
+		}()
 
 		if err := mgr.GetFieldIndexer().IndexField(context.Background(), &workspacev1.Workspace{}, "status.runtime.nodeName", func(o client.Object) []string {
 			ws := o.(*workspacev1.Workspace)
@@ -135,6 +144,17 @@ var runCmd = &cobra.Command{
 			return
 		}
 
+		if err := mgr.GetFieldIndexer().IndexField(context.Background(), &corev1.Pod{}, "spec.nodeName", func(o client.Object) []string {
+			pod := o.(*corev1.Pod)
+			if pod.Spec.NodeName == "" {
+				return nil
+			}
+			return []string{pod.Spec.NodeName}
+		}); err != nil {
+			log.WithError(err).Fatal("unable to create pod indexer")
+			return
+		}
+
 		nsac, err := NewNodeScaledownAnnotationController(mgr.GetClient())
 		if err != nil {
 			log.WithError(err).Fatal("unable to create node scaledown annotation controller")
@@ -153,10 +173,6 @@ var runCmd = &cobra.Command{
 		if err != nil {
 			log.WithError(err).Fatal("couldn't properly clean up node scaledown annotation controller")
 		}
-
-		metrics.Registry.MustRegister(NodeLabelerCounterVec)
-		metrics.Registry.MustRegister(NodeLabelerTimeHistVec)
-
 		err = mgr.AddHealthzCheck("healthz", healthz.Ping)
 		if err != nil {
 			log.WithError(err).Fatal("unable to set up health check")
@@ -208,96 +224,170 @@ func (r *PodReconciler) Reconcile(ctx context.Context, req reconcile.Request) (r
 		return reconcile.Result{RequeueAfter: defaultRequeueTime}, nil
 	}
 
+	var taintKey string
+	switch {
+	case strings.HasPrefix(pod.Name, registryFacade):
+		taintKey = registryFacadeTaintKey
+	case strings.HasPrefix(pod.Name, wsDaemon):
+		taintKey = wsDaemonTaintKey
+	default:
+		// nothing to do
+		return reconcile.Result{}, nil
+	}
+
+	healthy, err := checkPodHealth(pod)
+	if err != nil {
+		log.WithError(err).Error("checking pod health")
+		return reconcile.Result{RequeueAfter: defaultRequeueTime}, nil
+	}
+
+	var node corev1.Node
+	err = r.Get(ctx, types.NamespacedName{Name: nodeName}, &node)
+	if err != nil {
+		return reconcile.Result{}, fmt.Errorf("obtaining node %s: %w", nodeName, err)
+	}
+
+	if isNodeTaintExists(taintKey, node) != healthy {
+		// nothing to do, the taint already exists and is in the desired state.
+		return reconcile.Result{}, nil
+	}
+
+	err = updateNodeTaint(taintKey, !healthy, nodeName, r)
+	if err != nil {
+		log.WithError(err).Error("updating node taint")
+		return reconcile.Result{RequeueAfter: defaultRequeueTime}, nil
+	}
+
+	return reconcile.Result{}, nil
+}
+
+func checkPodHealth(pod corev1.Pod) (bool, error) {
 	var (
 		ipAddress string
 		port      string
-		taintKey  string
 	)
-
 	switch {
 	case strings.HasPrefix(pod.Name, registryFacade):
-		taintKey = registryFacadeTaintKey
 		ipAddress = pod.Status.HostIP
 		port = strconv.Itoa(registryFacadePort)
 	case strings.HasPrefix(pod.Name, wsDaemon):
-		taintKey = wsDaemonTaintKey
 		ipAddress = pod.Status.PodIP
 		port = strconv.Itoa(wsdaemonPort)
 	default:
 		// nothing to do
-		return reconcile.Result{}, nil
+		return true, nil
 	}
 
 	if !pod.ObjectMeta.DeletionTimestamp.IsZero() {
 		// the pod is being removed.
 		// add the taint to the node
-		time.Sleep(1 * time.Second)
-		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("adding node taint")
-			return reconcile.Result{RequeueAfter: defaultRequeueTime}, err
-		}
-
-		return reconcile.Result{}, err
+		return false, nil
 	}
 
 	if !IsPodReady(pod) {
 		// not ready. Wait until the next update.
-		return reconcile.Result{}, nil
+		return false, nil
 	}
 
-	var node corev1.Node
-	err = r.Get(ctx, types.NamespacedName{Name: nodeName}, &node)
+	err := checkTCPPortIsReachable(ipAddress, port)
 	if err != nil {
-		return reconcile.Result{}, fmt.Errorf("obtaining node %s: %w", nodeName, err)
+		log.WithField("host", ipAddress).WithField("port", port).WithField("pod", pod.Name).WithError(err).Error("checking if TCP port is open")
+		return false, nil
 	}
 
-	// Check if taint exists
-	taintExists := false
-	for _, taint := range node.Spec.Taints {
-		if taint.Key == taintKey {
-			taintExists = true
-			break
+	if strings.HasPrefix(pod.Name, registryFacade) {
+		err = checkRegistryFacade(ipAddress, port)
+		if err != nil {
+			log.WithError(err).Error("checking registry-facade")
+			return false, nil
 		}
 	}
 
-	if !taintExists {
-		// nothing to do, the taint doesn't exist.
-		return reconcile.Result{}, nil
+	return true, nil
+}
+
+type NodeReconciler struct {
+	client.Client
+}
+
+func (r *NodeReconciler) nodeFilter() predicate.Predicate {
+	return predicate.Funcs{
+		CreateFunc: func(e event.CreateEvent) bool {
+			node, ok := e.Object.(*corev1.Node)
+			if !ok {
+				return false
+			}
+			return isWorkspaceNode(*node)
+		},
+		UpdateFunc: func(e event.UpdateEvent) bool {
+			return false
+		},
+		DeleteFunc: func(e event.DeleteEvent) bool {
+			return false
+		},
 	}
+}
 
-	err = checkTCPPortIsReachable(ipAddress, port)
-	if err != nil {
-		log.WithField("host", ipAddress).WithField("port", port).WithField("pod", pod.Name).WithError(err).Error("checking if TCP port is open")
-		return reconcile.Result{RequeueAfter: defaultRequeueTime}, nil
+func (r *NodeReconciler) reconcileAll(ctx context.Context) error {
+	log.Info("start reconciling all nodes")
+
+	var nodes corev1.NodeList
+	if err := r.List(ctx, &nodes); err != nil {
+		return fmt.Errorf("failed to list nodes: %w", err)
 	}
 
-	if strings.HasPrefix(pod.Name, registryFacade) {
-		err = checkRegistryFacade(ipAddress, port)
-		if err != nil {
-			log.WithError(err).Error("checking registry-facade")
-			return reconcile.Result{RequeueAfter: defaultRequeueTime}, nil
+	for _, node := range nodes.Items {
+		if node.Labels == nil {
+			continue
+		}
+		if !isWorkspaceNode(node) {
+			continue
 		}
 
-		time.Sleep(1 * time.Second)
+		err := updateNodeLabel(node.Name, r.Client)
+		if err != nil {
+			log.WithError(err).WithField("node", node.Name).Error("failed to initialize labels on node")
+		}
+		r.Reconcile(ctx, reconcile.Request{NamespacedName: types.NamespacedName{Name: node.Name}})
 	}
 
-	err = updateNodeTaint(taintKey, false, nodeName, r)
+	log.Info("finished reconciling all nodes")
+	return nil
+}
+
+func (r *NodeReconciler) Reconcile(ctx context.Context, req reconcile.Request) (reconcile.Result, error) {
+	var node corev1.Node
+	err := r.Get(ctx, req.NamespacedName, &node)
 	if err != nil {
-		log.WithError(err).Error("removing node taint")
-		return reconcile.Result{}, fmt.Errorf("trying to remove the taint: %v", err)
+		if !errors.IsNotFound(err) {
+			log.WithError(err).Error("unable to fetch node")
+		}
+		return ctrl.Result{}, client.IgnoreNotFound(err)
+	}
+	var podList corev1.PodList
+	err = r.List(ctx, &podList, client.MatchingFields{
+		"spec.nodeName": node.Name,
+	})
+	if err != nil {
+		return reconcile.Result{}, fmt.Errorf("listing pods: %w", err)
+	}
+	isWsdaemonTaintExists := isNodeTaintExists(wsDaemonTaintKey, node)
+	isRegistryFacadeTaintExists := isNodeTaintExists(registryFacadeTaintKey, node)
+	isWsDaemonReady, isRegistryFacadeReady := false, false
+	for _, pod := range podList.Items {
+		if strings.HasPrefix(pod.Name, wsDaemon) {
+			isWsDaemonReady = IsPodReady(pod)
+		}
+		if strings.HasPrefix(pod.Name, registryFacade) {
+			isRegistryFacadeReady = IsPodReady(pod)
+		}
+	}
+	if isWsDaemonReady == isWsdaemonTaintExists {
+		updateNodeTaint(wsDaemonTaintKey, isWsDaemonReady, node.Name, r)
+	}
+	if isRegistryFacadeReady == isRegistryFacadeTaintExists {
+		updateNodeTaint(registryFacadeTaintKey, isRegistryFacadeReady, node.Name, r)
 	}
-
-	readyIn := time.Since(pod.Status.StartTime.Time)
-	NodeLabelerTimeHistVec.WithLabelValues(strings.Split(pod.Name, "-")[0]).Observe(readyIn.Seconds())
-	NodeLabelerCounterVec.WithLabelValues(strings.Split(pod.Name, "-")[0]).Inc()
-
 	return reconcile.Result{}, nil
 }
 
@@ -510,7 +600,10 @@ func updateNodeTaint(taintKey string, add bool, nodeName string, client client.C
 		var node corev1.Node
 		err := client.Get(ctx, types.NamespacedName{Name: nodeName}, &node)
 		if err != nil {
-			return err
+			if !errors.IsNotFound(err) {
+				return err
+			}
+			return nil
 		}
 
 		// Create or remove taint
@@ -554,6 +647,15 @@ func updateNodeTaint(taintKey string, add bool, nodeName string, client client.C
 	})
 }
 
+func isNodeTaintExists(taintKey string, node corev1.Node) bool {
+	for _, taint := range node.Spec.Taints {
+		if taint.Key == taintKey {
+			return true
+		}
+	}
+	return false
+}
+
 func checkTCPPortIsReachable(host string, port string) error {
 	conn, err := net.DialTimeout("tcp", net.JoinHostPort(host, port), 1*time.Second)
 	if err != nil {
@@ -611,31 +713,10 @@ func newDefaultTransport() *http.Transport {
 	}
 }
 
-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 isWorkspaceNode(node corev1.Node) bool {
+	_, isRegularWorkspaceNode := node.Labels[workspacesRegularLabel]
+	_, isHeadlessWorkspaceNode := node.Labels[workspacesHeadlessLabel]
+	return isRegularWorkspaceNode || isHeadlessWorkspaceNode
 }
 
 func updateNodeLabel(nodeName string, client client.Client) error {