From 9ea391c65a22dc9030de15d2996d1c7a45987ed5 Mon Sep 17 00:00:00 2001
From: Matt Terwilliger <matt@terwilligers.com>
Date: Fri, 29 Sep 2023 12:28:57 -0400
Subject: [PATCH] k8s: Support circular owner references (#6238)

Signed-off-by: Matthew Terwilliger <matt.terwilliger@reddit.com>
---
 internal/k8s/owner_fetcher.go      | 18 +++++++--
 internal/k8s/owner_fetcher_test.go | 65 ++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 4 deletions(-)

diff --git a/internal/k8s/owner_fetcher.go b/internal/k8s/owner_fetcher.go
index 635d62fec0..b9503d7705 100644
--- a/internal/k8s/owner_fetcher.go
+++ b/internal/k8s/owner_fetcher.go
@@ -6,6 +6,7 @@ import (
 	"strings"
 	"sync"
 
+	"golang.org/x/exp/slices"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -147,6 +148,10 @@ func (v OwnerFetcher) getOrCreatePromise(id types.UID) (*objectTreePromise, bool
 }
 
 func (v OwnerFetcher) OwnerTreeOfRef(ctx context.Context, ref v1.ObjectReference) (result ObjectRefTree, err error) {
+	return v.ownerTreeOfRefHelper(ctx, ref, nil)
+}
+
+func (v OwnerFetcher) ownerTreeOfRefHelper(ctx context.Context, ref v1.ObjectReference, path []types.UID) (result ObjectRefTree, err error) {
 	uid := ref.UID
 	if uid == "" {
 		return ObjectRefTree{}, fmt.Errorf("Can only get owners of deployed entities")
@@ -172,7 +177,7 @@ func (v OwnerFetcher) OwnerTreeOfRef(ctx context.Context, ref v1.ObjectReference
 		}
 		return ObjectRefTree{}, err
 	}
-	return v.ownerTreeOfHelper(ctx, ref, meta)
+	return v.ownerTreeOfHelper(ctx, ref, meta, path)
 }
 
 func (v OwnerFetcher) getMetaByReference(ctx context.Context, ref v1.ObjectReference) (metav1.Object, error) {
@@ -211,15 +216,20 @@ func (v OwnerFetcher) OwnerTreeOf(ctx context.Context, entity K8sEntity) (result
 	}()
 
 	ref := entity.ToObjectReference()
-	return v.ownerTreeOfHelper(ctx, ref, meta)
+	return v.ownerTreeOfHelper(ctx, ref, meta, nil)
 }
 
-func (v OwnerFetcher) ownerTreeOfHelper(ctx context.Context, ref v1.ObjectReference, meta metav1.Object) (ObjectRefTree, error) {
+func (v OwnerFetcher) ownerTreeOfHelper(ctx context.Context, ref v1.ObjectReference, meta metav1.Object, path []types.UID) (ObjectRefTree, error) {
 	tree := ObjectRefTree{Ref: ref, CreationTimestamp: meta.GetCreationTimestamp()}
 	owners := meta.GetOwnerReferences()
 	for _, owner := range owners {
+		// TODO: Owner references can also exist at cluster scope, for which this incorrectly propagates the parent ref's Namespace.
 		ownerRef := OwnerRefToObjectRef(owner, meta.GetNamespace())
-		ownerTree, err := v.OwnerTreeOfRef(ctx, ownerRef)
+		if slices.Contains(path, owner.UID) {
+			// break circular dependencies
+			continue
+		}
+		ownerTree, err := v.ownerTreeOfRefHelper(ctx, ownerRef, append(path, ref.UID))
 		if err != nil {
 			return ObjectRefTree{}, err
 		}
diff --git a/internal/k8s/owner_fetcher_test.go b/internal/k8s/owner_fetcher_test.go
index dc15f1ffc3..01f5e919be 100644
--- a/internal/k8s/owner_fetcher_test.go
+++ b/internal/k8s/owner_fetcher_test.go
@@ -79,6 +79,21 @@ func TestOwnerFetcherParallelism(t *testing.T) {
 	assert.Equal(t, 1, kCli.getByReferenceCallCount)
 }
 
+func TestCircular(t *testing.T) {
+	kCli := NewFakeK8sClient(t)
+	kCli.listReturnsEmpty = true
+	ov := NewOwnerFetcher(context.Background(), kCli)
+
+	pod1, pod2, pod3 := fakeCircularReference()
+	kCli.Inject(NewK8sEntity(pod2), NewK8sEntity(pod3))
+
+	tree, err := ov.OwnerTreeOf(context.Background(), NewK8sEntity(pod1))
+	assert.NoError(t, err)
+	assert.Equal(t, `Pod:pod-a
+  Pod:pod-b
+    Pod:pod-c`, tree.String())
+}
+
 func fakeOneParentChain() (*v1.Pod, *appsv1.ReplicaSet) {
 	pod := &v1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
@@ -145,3 +160,53 @@ func fakeTwoParentChain() (*v1.Pod, *appsv1.ReplicaSet, *appsv1.Deployment) {
 	}
 	return pod, rs, dep
 }
+
+func fakeCircularReference() (*v1.Pod, *v1.Pod, *v1.Pod) {
+	pod1 := &v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "pod-a",
+			UID:       "pod-a-uid",
+			Namespace: "default",
+			OwnerReferences: []metav1.OwnerReference{
+				{
+					APIVersion: "v1",
+					Kind:       "Pod",
+					Name:       "pod-b",
+					UID:        "pod-b-uid",
+				},
+			},
+		},
+	}
+	pod2 := &v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "pod-b",
+			UID:       "pod-b-uid",
+			Namespace: "default",
+			OwnerReferences: []metav1.OwnerReference{
+				{
+					APIVersion: "v1",
+					Kind:       "Pod",
+					Name:       "pod-c",
+					UID:        "pod-c-uid",
+				},
+			},
+		},
+	}
+	pod3 := &v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "pod-c",
+			UID:       "pod-c-uid",
+			Namespace: "default",
+			OwnerReferences: []metav1.OwnerReference{
+				{
+					APIVersion: "v1",
+					Kind:       "Pod",
+					Name:       "pod-a",
+					UID:        "pod-a-uid",
+				},
+			},
+		},
+	}
+
+	return pod1, pod2, pod3
+}