From e8592c3827533e9c3b1a3e0dafee0e8c9b9b19ca Mon Sep 17 00:00:00 2001
From: fabriziopandini <fpandini@vmware.com>
Date: Tue, 14 Dec 2021 21:41:17 +0100
Subject: [PATCH] Prevent orphaned InfrastructureMachines

---
 controllers/machineset_controller.go      | 12 ++++++++++++
 controllers/machineset_controller_test.go | 22 ++++++++++++++++++----
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/controllers/machineset_controller.go b/controllers/machineset_controller.go
index 9e075b9dfe71..31f41d5e2e84 100644
--- a/controllers/machineset_controller.go
+++ b/controllers/machineset_controller.go
@@ -376,6 +376,12 @@ func (r *MachineSetReconciler) syncReplicas(ctx context.Context, ms *clusterv1.M
 					ClusterName: machine.Spec.ClusterName,
 					Labels:      machine.Labels,
 					Annotations: machine.Annotations,
+					OwnerRef: &metav1.OwnerReference{
+						APIVersion: clusterv1.GroupVersion.String(),
+						Kind:       "MachineSet",
+						Name:       ms.Name,
+						UID:        ms.UID,
+					},
 				})
 				if err != nil {
 					conditions.MarkFalse(ms, clusterv1.MachinesCreatedCondition, clusterv1.BootstrapTemplateCloningFailedReason, clusterv1.ConditionSeverityError, err.Error())
@@ -391,6 +397,12 @@ func (r *MachineSetReconciler) syncReplicas(ctx context.Context, ms *clusterv1.M
 				ClusterName: machine.Spec.ClusterName,
 				Labels:      machine.Labels,
 				Annotations: machine.Annotations,
+				OwnerRef: &metav1.OwnerReference{
+					APIVersion: clusterv1.GroupVersion.String(),
+					Kind:       "MachineSet",
+					Name:       ms.Name,
+					UID:        ms.UID,
+				},
 			})
 			if err != nil {
 				conditions.MarkFalse(ms, clusterv1.MachinesCreatedCondition, clusterv1.InfrastructureTemplateCloningFailedReason, clusterv1.ConditionSeverityError, err.Error())
diff --git a/controllers/machineset_controller_test.go b/controllers/machineset_controller_test.go
index 69e9eb97ba8d..fa67c972ad94 100644
--- a/controllers/machineset_controller_test.go
+++ b/controllers/machineset_controller_test.go
@@ -228,6 +228,13 @@ func TestMachineSetReconciler(t *testing.T) {
 			g.Expect(im.GetAnnotations()).To(HaveKeyWithValue("annotation-1", "true"), "have annotations of MachineTemplate applied")
 			g.Expect(im.GetAnnotations()).To(HaveKeyWithValue("precedence", "MachineSet"), "the annotations from the MachineSpec template to overwrite the infrastructure template ones")
 			g.Expect(im.GetLabels()).To(HaveKeyWithValue("label-1", "true"), "have labels of MachineTemplate applied")
+			hasMSOwnerRef := false
+			for _, o := range im.GetOwnerReferences() {
+				if o.Kind == machineSetKind.Kind {
+					hasMSOwnerRef = true
+				}
+			}
+			g.Expect(hasMSOwnerRef).To(BeTrue(), "have ownerRef to MachineSet")
 		}
 
 		t.Log("Creating a BootstrapConfig for each Machine")
@@ -244,6 +251,13 @@ func TestMachineSetReconciler(t *testing.T) {
 			g.Expect(im.GetAnnotations()).To(HaveKeyWithValue("annotation-1", "true"), "have annotations of MachineTemplate applied")
 			g.Expect(im.GetAnnotations()).To(HaveKeyWithValue("precedence", "MachineSet"), "the annotations from the MachineSpec template to overwrite the bootstrap config template ones")
 			g.Expect(im.GetLabels()).To(HaveKeyWithValue("label-1", "true"), "have labels of MachineTemplate applied")
+			hasMSOwnerRef := false
+			for _, o := range im.GetOwnerReferences() {
+				if o.Kind == machineSetKind.Kind {
+					hasMSOwnerRef = true
+				}
+			}
+			g.Expect(hasMSOwnerRef).To(BeTrue(), "have ownerRef to MachineSet")
 		}
 
 		// Set the infrastructure reference as ready.
@@ -527,7 +541,7 @@ func TestMachineSetToMachines(t *testing.T) {
 			OwnerReferences: []metav1.OwnerReference{
 				{
 					Name:       "Owner",
-					Kind:       "MachineSet",
+					Kind:       machineSetKind.Kind,
 					Controller: &controller,
 				},
 			},
@@ -608,7 +622,7 @@ func TestShouldExcludeMachine(t *testing.T) {
 					OwnerReferences: []metav1.OwnerReference{
 						{
 							Name:       "Owner",
-							Kind:       "MachineSet",
+							Kind:       machineSetKind.Kind,
 							Controller: &controller,
 							UID:        "not-1",
 						},
@@ -628,7 +642,7 @@ func TestShouldExcludeMachine(t *testing.T) {
 					OwnerReferences: []metav1.OwnerReference{
 						{
 							Name:       "Owner",
-							Kind:       "MachineSet",
+							Kind:       machineSetKind.Kind,
 							Controller: &controller,
 							UID:        "1",
 						},
@@ -709,7 +723,7 @@ func TestAdoptOrphan(t *testing.T) {
 			expected: []metav1.OwnerReference{
 				{
 					APIVersion:         clusterv1.GroupVersion.String(),
-					Kind:               "MachineSet",
+					Kind:               machineSetKind.Kind,
 					Name:               "adoptOrphanMachine",
 					UID:                "",
 					Controller:         &controller,