Apply labels to Machines if they're part of a MD or MS

Author: tahsinrahman

api/v1alpha3/machine_types.go

@@ -31,6 +31,12 @@ const (
 
 	// ExcludeNodeDrainingAnnotation annotation explicitly skips node draining if set
 	ExcludeNodeDrainingAnnotation = "machine.cluster.x-k8s.io.io/exclude-node-draining"
+
+	// MachineSetLabelName is the label set on machines if they're controlled by MachineSet
+	MachineSetLabelName = "cluster.x-k8s.io/set-name"
+
+	// MachineDeploymentLabelName is the label set on machines if they're controlled by MachineDeployment
+	MachineDeploymentLabelName = "cluster.x-k8s.io/deployment-name"
 )
 
 // ANCHOR: MachineSpec

controllers/machine_controller.go

@@ -157,6 +157,7 @@ func (r *MachineReconciler) reconcile(ctx context.Context, cluster *clusterv1.Cl
 
 	// Call the inner reconciliation methods.
 	reconciliationErrors := []error{
+		r.reconcileLables(ctx, m),
 		r.reconcileBootstrap(ctx, m),
 		r.reconcileInfrastructure(ctx, m),
 		r.reconcileNodeRef(ctx, cluster, m),
@@ -181,6 +182,42 @@ func (r *MachineReconciler) reconcile(ctx context.Context, cluster *clusterv1.Cl
 	return res, kerrors.NewAggregate(errs)
 }
 
+func (r *MachineReconciler) reconcileLables(ctx context.Context, m *clusterv1.Machine) error {
+	logger := r.Log.WithValues("machine", m.Name, "namespace", m.Namespace)
+
+	if m.Labels == nil {
+		m.Labels = make(map[string]string)
+	}
+
+	// return early if machine already has both the labels for machineSet and machineDeployment controller
+	if m.Labels[clusterv1.MachineDeploymentLabelName] != "" && m.Labels[clusterv1.MachineSetLabelName] != "" {
+		return nil
+	}
+
+	machineControllerRef := metav1.GetControllerOf(m)
+	// return early if machine is not controlled by any machineSet
+	if machineControllerRef == nil {
+		return nil
+	}
+
+	ms := &clusterv1.MachineSet{}
+	err := r.Client.Get(ctx, client.ObjectKey{
+		Namespace: m.Namespace,
+		Name:      machineControllerRef.Name,
+	}, ms)
+	if err != nil {
+		logger.Error(err, "failed to get owner machineset", "machineset", machineControllerRef.Name)
+		return err
+	}
+	m.Labels[clusterv1.MachineSetLabelName] = machineControllerRef.Name
+
+	if machineSetControllerRef := metav1.GetControllerOf(ms); machineSetControllerRef != nil {
+		m.Labels[clusterv1.MachineDeploymentLabelName] = machineSetControllerRef.Name
+	}
+
+	return nil
+}
+
 func (r *MachineReconciler) reconcileDelete(ctx context.Context, cluster *clusterv1.Cluster, m *clusterv1.Machine) (ctrl.Result, error) {
 	logger := r.Log.WithValues("machine", m.Name, "namespace", m.Namespace)
 	logger = logger.WithValues("cluster", cluster.Name)

controllers/machine_controller_test.go

@@ -17,6 +17,8 @@ limitations under the License.
 package controllers
 
 import (
+	"context"
+	"reflect"
 	"testing"
 
 	. "github.com/onsi/gomega"
@@ -26,6 +28,7 @@ import (
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/kubernetes/scheme"
+	"k8s.io/klog/klogr"
 	"k8s.io/utils/pointer"
 	clusterv1 "sigs.k8s.io/cluster-api/api/v1alpha3"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -553,3 +556,125 @@ func TestRemoveMachineFinalizerAfterDeleteReconcile(t *testing.T) {
 	Expect(mr.Client.Get(ctx, key, m)).ToNot(HaveOccurred())
 	Expect(m.ObjectMeta.Finalizers).To(Equal([]string{metav1.FinalizerDeleteDependents}))
 }
+
+func TestMachineReconciler_reconcileLables(t *testing.T) {
+	RegisterTestingT(t)
+
+	err := clusterv1.AddToScheme(scheme.Scheme)
+	Expect(err).NotTo(HaveOccurred())
+
+	tests := []struct {
+		name           string
+		m              *clusterv1.Machine
+		ms             *clusterv1.MachineSet
+		md             *clusterv1.MachineDeployment
+		expectedLabels map[string]string
+		wantErr        bool
+	}{
+		{
+			name: "no controller, machine should not have labels",
+			m: &clusterv1.Machine{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-machine",
+					Namespace: corev1.NamespaceDefault,
+				},
+			},
+			expectedLabels: map[string]string{},
+			wantErr:        false,
+		},
+		{
+			name: "controlled by machineset",
+			m: &clusterv1.Machine{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-machine",
+					Namespace: corev1.NamespaceDefault,
+					OwnerReferences: []metav1.OwnerReference{
+						{
+							APIVersion:         clusterv1.GroupVersion.String(),
+							Kind:               machineSetKind.String(),
+							Name:               "test-ms",
+							Controller:         pointer.BoolPtr(true),
+							BlockOwnerDeletion: nil,
+						},
+					},
+				},
+			},
+			ms: &clusterv1.MachineSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-ms",
+					Namespace: corev1.NamespaceDefault,
+				},
+			},
+			expectedLabels: map[string]string{
+				clusterv1.MachineSetLabelName: "test-ms",
+			},
+			wantErr: false,
+		},
+		{
+			name: "controlled by machinedeployment",
+			m: &clusterv1.Machine{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-machine",
+					Namespace: corev1.NamespaceDefault,
+					OwnerReferences: []metav1.OwnerReference{
+						{
+							APIVersion:         clusterv1.GroupVersion.String(),
+							Kind:               machineSetKind.String(),
+							Name:               "test-ms",
+							Controller:         pointer.BoolPtr(true),
+							BlockOwnerDeletion: nil,
+						},
+					},
+				},
+			},
+			ms: &clusterv1.MachineSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-ms",
+					Namespace: corev1.NamespaceDefault,
+					OwnerReferences: []metav1.OwnerReference{
+						{
+							APIVersion:         clusterv1.GroupVersion.String(),
+							Kind:               machineDeploymentKind.String(),
+							Name:               "test-md",
+							Controller:         pointer.BoolPtr(true),
+							BlockOwnerDeletion: nil,
+						},
+					},
+				},
+			},
+			md: &clusterv1.MachineDeployment{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-md",
+					Namespace: corev1.NamespaceDefault,
+				},
+			},
+			expectedLabels: map[string]string{
+				clusterv1.MachineSetLabelName:        "test-ms",
+				clusterv1.MachineDeploymentLabelName: "test-md",
+			},
+			wantErr: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := fake.NewFakeClientWithScheme(scheme.Scheme, tt.m)
+			if tt.ms != nil {
+				c = fake.NewFakeClientWithScheme(scheme.Scheme, tt.m, tt.ms)
+			}
+			if tt.md != nil {
+				c = fake.NewFakeClientWithScheme(scheme.Scheme, tt.m, tt.ms, tt.md)
+			}
+
+			r := &MachineReconciler{
+				Log:    klogr.New(),
+				Client: c,
+			}
+			if err := r.reconcileLables(context.Background(), tt.m); (err != nil) != tt.wantErr {
+				t.Errorf("reconcileLables() error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if !reflect.DeepEqual(tt.m.Labels, tt.expectedLabels) {
+				t.Errorf("expected labels %v, got %v", tt.expectedLabels, tt.m.Labels)
+			}
+		})
+	}
+}