From c13c2c2b29764185243a5eeaaf7da7f7d946a421 Mon Sep 17 00:00:00 2001
From: Andrew Lavery <laverya@umich.edu>
Date: Fri, 20 Oct 2023 10:28:25 -0600
Subject: [PATCH] node role labels (#4093)

* node role labels

* handle having no labels on the first node

* f

* include a prefix on the label

* = not :
---
 pkg/helmvm/helmvm_node.go | 26 ++++++++++++++++++++-----
 pkg/helmvm/node_join.go   | 41 +++++++++++++++++++++++++++++++++++++++
 pkg/helmvm/types/types.go |  3 +++
 3 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/pkg/helmvm/helmvm_node.go b/pkg/helmvm/helmvm_node.go
index 079086f983..a9b17f497e 100644
--- a/pkg/helmvm/helmvm_node.go
+++ b/pkg/helmvm/helmvm_node.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"math"
 	"strconv"
+	"strings"
 
 	"github.com/replicatedhq/kots/pkg/helmvm/types"
 	"github.com/replicatedhq/kots/pkg/k8sutil"
@@ -136,16 +137,31 @@ func nodeMetrics(ctx context.Context, client kubernetes.Interface, metricsClient
 	}, nil
 }
 
+// nodeRolesFromLabels parses a map of k8s node labels, and returns the roles of the node
 func nodeRolesFromLabels(labels map[string]string) []string {
 	toReturn := []string{}
 
-	// detect if this is a controller node from the k8s labels
-	if val, ok := labels["node-role.kubernetes.io/control-plane"]; ok && val == "true" {
-		toReturn = append(toReturn, "controller")
+	numRolesStr, ok := labels[types.EMBEDDED_CLUSTER_ROLE_LABEL]
+	if !ok {
+		// the first node will not initially have a role label, but is a 'controller'
+		if val, ok := labels["node-role.kubernetes.io/control-plane"]; ok && val == "true" {
+			return []string{"controller"}
+		}
+		return nil
 	}
+	numRoles, err := strconv.Atoi(strings.TrimPrefix(numRolesStr, "total-"))
+	if err != nil {
+		fmt.Printf("failed to parse role label %q: %s", numRolesStr, err.Error())
 
-	if len(toReturn) == 0 {
-		toReturn = append(toReturn, "worker")
+		return nil
+	}
+
+	for i := 0; i < numRoles; i++ {
+		roleLabel, ok := labels[fmt.Sprintf("%s-%d", types.EMBEDDED_CLUSTER_ROLE_LABEL, i)]
+		if !ok {
+			fmt.Printf("failed to find role label %d", i)
+		}
+		toReturn = append(toReturn, roleLabel)
 	}
 
 	return toReturn
diff --git a/pkg/helmvm/node_join.go b/pkg/helmvm/node_join.go
index ea739c4b54..4a4bb6c068 100644
--- a/pkg/helmvm/node_join.go
+++ b/pkg/helmvm/node_join.go
@@ -8,6 +8,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/replicatedhq/kots/pkg/helmvm/types"
 	corev1 "k8s.io/api/core/v1"
 	kuberneteserrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -252,6 +253,12 @@ func GenerateK0sJoinCommand(ctx context.Context, client kubernetes.Interface, ro
 		cmd = append(cmd, "--enable-worker")
 	}
 
+	labels, err := getRolesNodeLabels(ctx, client, roles)
+	if err != nil {
+		return "", fmt.Errorf("failed to get role labels: %w", err)
+	}
+	cmd = append(cmd, "--labels", labels)
+
 	return strings.Join(cmd, " "), nil
 }
 
@@ -296,3 +303,37 @@ func getControllerNodeIP(ctx context.Context, client kubernetes.Interface) (stri
 
 	return "", fmt.Errorf("failed to find healthy controller node")
 }
+
+func getRolesNodeLabels(ctx context.Context, client kubernetes.Interface, roles []string) (string, error) {
+	roleLabels := getRoleListLabels(roles)
+
+	for _, role := range roles {
+		labels, err := getRoleNodeLabels(ctx, client, role)
+		if err != nil {
+			return "", fmt.Errorf("failed to get node labels for role %s: %w", role, err)
+		}
+		roleLabels = append(roleLabels, labels...)
+	}
+
+	return strings.Join(roleLabels, ","), nil
+}
+
+// TODO: look up role in cluster config, apply additional labels based on role
+func getRoleNodeLabels(ctx context.Context, client kubernetes.Interface, role string) ([]string, error) {
+	toReturn := []string{}
+
+	return toReturn, nil
+}
+
+// getRoleListLabels returns the labels needed to identify the roles of this node in the future
+// one label will be the number of roles, and then deterministic label names will be used to store the role names
+func getRoleListLabels(roles []string) []string {
+	toReturn := []string{}
+	toReturn = append(toReturn, fmt.Sprintf("%s=total-%d", types.EMBEDDED_CLUSTER_ROLE_LABEL, len(roles)))
+
+	for idx, role := range roles {
+		toReturn = append(toReturn, fmt.Sprintf("%s-%d=%s", types.EMBEDDED_CLUSTER_ROLE_LABEL, idx, role))
+	}
+
+	return toReturn
+}
diff --git a/pkg/helmvm/types/types.go b/pkg/helmvm/types/types.go
index f177df2b37..10bf390368 100644
--- a/pkg/helmvm/types/types.go
+++ b/pkg/helmvm/types/types.go
@@ -1,5 +1,8 @@
 package types
 
+const EMBEDDED_CLUSTER_LABEL = "kots.io/embedded-cluster"
+const EMBEDDED_CLUSTER_ROLE_LABEL = EMBEDDED_CLUSTER_LABEL + "-role"
+
 type HelmVMNodes struct {
 	Nodes           []Node `json:"nodes"`
 	HA              bool   `json:"ha"`