✨ Add controllerutil.EnsureOwnerReference

Signed-off-by: Vince Prignano <vincepri@vmware.com>
kubernetes-sigs · Feb 24, 2020 · 8fc450e · 8fc450e
1 parent 525a922
commit 8fc450e
Show file tree

Hide file tree

Showing 2 changed files with 152 additions and 27 deletions.
diff --git a/pkg/controller/controllerutil/controllerutil.go b/pkg/controller/controllerutil/controllerutil.go
@@ -55,47 +55,95 @@ func newAlreadyOwnedError(Object metav1.Object, Owner metav1.OwnerReference) *Al
 // Since only one OwnerReference can be a controller, it returns an error if
 // there is another OwnerReference with Controller flag set.
 func SetControllerReference(owner, controlled metav1.Object, scheme *runtime.Scheme) error {
+	// Return early with an error if the object is already controlled.
+	if existing := metav1.GetControllerOf(controlled); existing != nil {
+		return newAlreadyOwnedError(controlled, *existing)
+	}
+
+	// Validate the owner.
 	ro, ok := owner.(runtime.Object)
 	if !ok {
 		return fmt.Errorf("%T is not a runtime.Object, cannot call SetControllerReference", owner)
 	}
-
-	ownerNs := owner.GetNamespace()
-	if ownerNs != "" {
-		objNs := controlled.GetNamespace()
-		if objNs == "" {
-			return fmt.Errorf("cluster-scoped resource must not have a namespace-scoped owner, owner's namespace %s", ownerNs)
-		}
-		if ownerNs != objNs {
-			return fmt.Errorf("cross-namespace owner references are disallowed, owner's namespace %s, obj's namespace %s", owner.GetNamespace(), controlled.GetNamespace())
-		}
+	if err := validateOwner(owner, controlled); err != nil {
+		return err
 	}
 
+	// Create a new controller ref.
 	gvk, err := apiutil.GVKForObject(ro, scheme)
 	if err != nil {
 		return err
 	}
-
-	// Create a new ref
 	ref := *metav1.NewControllerRef(owner, schema.GroupVersionKind{Group: gvk.Group, Version: gvk.Version, Kind: gvk.Kind})
 
-	existingRefs := controlled.GetOwnerReferences()
-	fi := -1
-	for i, r := range existingRefs {
-		if referSameObject(ref, r) {
-			fi = i
-		} else if r.Controller != nil && *r.Controller {
-			return newAlreadyOwnedError(controlled, r)
-		}
+	// Update owner references and return.
+	upsertOwnerRef(ref, controlled)
+	return nil
+}
+
+// EnsureOwnerReference is a helper method to make sure the given object contains
+// an object reference to the object provided.
+// If a reference already exists, it'll be overwritten with the newly provided version.
+func EnsureOwnerReference(owner, object metav1.Object, scheme *runtime.Scheme) error {
+	// Validate the owner.
+	ro, ok := owner.(runtime.Object)
+	if !ok {
+		return fmt.Errorf("%T is not a runtime.Object, cannot call SetControllerReference", owner)
+	}
+	if err := validateOwner(owner, object); err != nil {
+		return err
 	}
-	if fi == -1 {
-		existingRefs = append(existingRefs, ref)
+
+	// Create a new owner ref.
+	gvk, err := apiutil.GVKForObject(ro, scheme)
+	if err != nil {
+		return err
+	}
+	ref := metav1.OwnerReference{
+		APIVersion: gvk.GroupVersion().String(),
+		Kind:       gvk.Kind,
+		UID:        owner.GetUID(),
+		Name:       owner.GetName(),
+	}
+
+	// Update owner references and return.
+	upsertOwnerRef(ref, object)
+	return nil
+
+}
+
+func upsertOwnerRef(ref metav1.OwnerReference, object metav1.Object) {
+	owners := object.GetOwnerReferences()
+	idx := indexOwnerRef(owners, ref)
+	if idx == -1 {
+		owners = append(owners, ref)
 	} else {
-		existingRefs[fi] = ref
+		owners[idx] = ref
+	}
+	object.SetOwnerReferences(owners)
+}
+
+// indexOwnerRef returns the index of the owner reference in the slice if found, or -1.
+func indexOwnerRef(ownerReferences []metav1.OwnerReference, ref metav1.OwnerReference) int {
+	for index, r := range ownerReferences {
+		if referSameObject(r, ref) {
+			return index
+		}
 	}
+	return -1
+}
 
-	// Update owner references
-	controlled.SetOwnerReferences(existingRefs)
+func validateOwner(owner, object metav1.Object) error {
+	ownerNs := owner.GetNamespace()
+	if ownerNs != "" {
+		objNs := object.GetNamespace()
+		if objNs == "" {
+			return fmt.Errorf("cluster-scoped resource must not have a namespace-scoped owner, owner's namespace %s", ownerNs)
+		}
+		if ownerNs != objNs {
+			return fmt.Errorf("cross-namespace owner references are disallowed, owner's namespace %s, obj's namespace %s", owner.GetNamespace(), object.GetNamespace())
+		}
+	}
 	return nil
 }
 

diff --git a/pkg/controller/controllerutil/controllerutil_test.go b/pkg/controller/controllerutil/controllerutil_test.go
@@ -21,8 +21,6 @@ import (
 	"fmt"
 	"math/rand"
 
-	"sigs.k8s.io/controller-runtime/pkg/client"
-
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 	appsv1 "k8s.io/api/apps/v1"
@@ -32,10 +30,89 @@ import (
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/kubernetes/scheme"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
 )
 
 var _ = Describe("Controllerutil", func() {
+	Describe("EnsureOwnerReference", func() {
+		It("should set ownerRef on an empty list", func() {
+			rs := &appsv1.ReplicaSet{}
+			dep := &extensionsv1beta1.Deployment{
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", UID: "foo-uid"},
+			}
+			Expect(controllerutil.EnsureOwnerReference(dep, rs, scheme.Scheme)).ToNot(HaveOccurred())
+			t := false
+			Expect(rs.OwnerReferences).To(ConsistOf(metav1.OwnerReference{
+				Name:               "foo",
+				Kind:               "Deployment",
+				APIVersion:         "extensions/v1beta1",
+				UID:                "foo-uid",
+				Controller:         &t,
+				BlockOwnerDeletion: &t,
+			}))
+		})
+
+		It("should not duplicate owner references", func() {
+			rs := &appsv1.ReplicaSet{
+				ObjectMeta: metav1.ObjectMeta{
+					OwnerReferences: []metav1.OwnerReference{
+						{
+							Name:       "foo",
+							Kind:       "Deployment",
+							APIVersion: "extensions/v1beta1",
+							UID:        "foo-uid",
+						},
+					},
+				},
+			}
+			dep := &extensionsv1beta1.Deployment{
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", UID: "foo-uid"},
+			}
+
+			Expect(controllerutil.EnsureOwnerReference(dep, rs, scheme.Scheme)).ToNot(HaveOccurred())
+			t := false
+			Expect(rs.OwnerReferences).To(ConsistOf(metav1.OwnerReference{
+				Name:               "foo",
+				Kind:               "Deployment",
+				APIVersion:         "extensions/v1beta1",
+				UID:                "foo-uid",
+				Controller:         &t,
+				BlockOwnerDeletion: &t,
+			}))
+		})
+
+		It("should update the APIVersion if duplicate", func() {
+			rs := &appsv1.ReplicaSet{
+				ObjectMeta: metav1.ObjectMeta{
+					OwnerReferences: []metav1.OwnerReference{
+						{
+							Name:       "foo",
+							Kind:       "Deployment",
+							APIVersion: "extensions/v1alpha1",
+							UID:        "foo-uid",
+						},
+					},
+				},
+			}
+			dep := &extensionsv1beta1.Deployment{
+				ObjectMeta: metav1.ObjectMeta{Name: "foo", UID: "foo-uid"},
+			}
+
+			Expect(controllerutil.EnsureOwnerReference(dep, rs, scheme.Scheme)).ToNot(HaveOccurred())
+			t := false
+			Expect(rs.OwnerReferences).To(ConsistOf(metav1.OwnerReference{
+				Name:               "foo",
+				Kind:               "Deployment",
+				APIVersion:         "extensions/v1beta1",
+				UID:                "foo-uid",
+				Controller:         &t,
+				BlockOwnerDeletion: &t,
+			}))
+
+		})
+	})
+
 	Describe("SetControllerReference", func() {
 		It("should set the OwnerReference if it can find the group version kind", func() {
 			rs := &appsv1.ReplicaSet{}