bimapper: allow to untrack links and support reference or id (#18451)

internal/catalog/internal/mappers/failovermapper/failover_mapper.go

@@ -42,7 +42,11 @@ func (m *Mapper) TrackFailover(failover *resource.DecodedResource[pbcatalog.Fail
 }
 
 func (m *Mapper) trackFailover(failover *pbresource.ID, services []*pbresource.Reference) {
-	m.b.TrackItem(failover, services)
+	var servicesAsIDsOrRefs []resource.ReferenceOrID
+	for _, s := range services {
+		servicesAsIDsOrRefs = append(servicesAsIDsOrRefs, s)
+	}
+	m.b.TrackItem(failover, servicesAsIDsOrRefs)
 }
 
 // UntrackFailover forgets the links inserted by TrackFailover for the provided

internal/resource/equality.go

@@ -20,7 +20,7 @@ func EqualType(a, b *pbresource.Type) bool {
 		a.Kind == b.Kind
 }
 
-// EqualType compares two resource tenancies for equality without reflection.
+// EqualTenancy compares two resource tenancies for equality without reflection.
 func EqualTenancy(a, b *pbresource.Tenancy) bool {
 	if a == b {
 		return true
@@ -35,7 +35,7 @@ func EqualTenancy(a, b *pbresource.Tenancy) bool {
 		a.Namespace == b.Namespace
 }
 
-// EqualType compares two resource IDs for equality without reflection.
+// EqualID compares two resource IDs for equality without reflection.
 func EqualID(a, b *pbresource.ID) bool {
 	if a == b {
 		return true

internal/resource/mappers/bimapper/bimapper.go

@@ -59,16 +59,34 @@ func (m *Mapper) IsEmpty() bool {
 
 // UntrackItem removes tracking for the provided item. The item type MUST match
 // the type configured for the item.
-func (m *Mapper) UntrackItem(item *pbresource.ID) {
+func (m *Mapper) UntrackItem(item resource.ReferenceOrID) {
-	if !resource.EqualType(item.Type, m.itemType) {
+	if !resource.EqualType(item.GetType(), m.itemType) {
 		panic(fmt.Sprintf("expected item type %q got %q",
 			resource.TypeToString(m.itemType),
-			resource.TypeToString(item.Type),
+			resource.TypeToString(item.GetType()),
 		))
 	}
 	m.untrackItem(resource.NewReferenceKey(item))
 }
 
+// UntrackLink removes tracking for the provided link. The link type MUST match
+// the type configured for the link.
+func (m *Mapper) UntrackLink(link resource.ReferenceOrID) {
+	if !resource.EqualType(link.GetType(), m.linkType) {
+		panic(fmt.Sprintf("expected link type %q got %q",
+			resource.TypeToString(m.linkType),
+			resource.TypeToString(link.GetType()),
+		))
+	}
+	m.untrackLink(resource.NewReferenceKey(link))
+}
+
+func (m *Mapper) untrackLink(link resource.ReferenceKey) {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+	m.removeLinkLocked(link)
+}
+
 func (m *Mapper) untrackItem(item resource.ReferenceKey) {
 	m.lock.Lock()
 	defer m.lock.Unlock()
@@ -77,20 +95,20 @@ func (m *Mapper) untrackItem(item resource.ReferenceKey) {
 
 // TrackItem adds tracking for the provided item. The item and link types MUST
 // match the types configured for the items and links.
-func (m *Mapper) TrackItem(item *pbresource.ID, links []*pbresource.Reference) {
+func (m *Mapper) TrackItem(item resource.ReferenceOrID, links []resource.ReferenceOrID) {
-	if !resource.EqualType(item.Type, m.itemType) {
+	if !resource.EqualType(item.GetType(), m.itemType) {
 		panic(fmt.Sprintf("expected item type %q got %q",
 			resource.TypeToString(m.itemType),
-			resource.TypeToString(item.Type),
+			resource.TypeToString(item.GetType()),
 		))
 	}
 
 	linksAsKeys := make([]resource.ReferenceKey, 0, len(links))
 	for _, link := range links {
-		if !resource.EqualType(link.Type, m.linkType) {
+		if !resource.EqualType(link.GetType(), m.linkType) {
 			panic(fmt.Sprintf("expected link type %q got %q",
 				resource.TypeToString(m.linkType),
-				resource.TypeToString(link.Type),
+				resource.TypeToString(link.GetType()),
 			))
 		}
 		linksAsKeys = append(linksAsKeys, resource.NewReferenceKey(link))
@@ -118,6 +136,16 @@ func (m *Mapper) removeItemLocked(item resource.ReferenceKey) {
 	delete(m.itemToLink, item)
 }
 
+func (m *Mapper) removeLinkLocked(link resource.ReferenceKey) {
+	for item := range m.linkToItem[link] {
+		delete(m.itemToLink[item], link)
+		if len(m.itemToLink[item]) == 0 {
+			delete(m.itemToLink, item)
+		}
+	}
+	delete(m.linkToItem, link)
+}
+
 // you must hold the lock before calling this function
 func (m *Mapper) addItemLocked(item resource.ReferenceKey, links []resource.ReferenceKey) {
 	if m.itemToLink[item] == nil {
@@ -134,7 +162,13 @@ func (m *Mapper) addItemLocked(item resource.ReferenceKey, links []resource.Refe
 }
 
 // LinksForItem returns references to links related to the requested item.
+// Deprecated: use LinksRefs
 func (m *Mapper) LinksForItem(item *pbresource.ID) []*pbresource.Reference {
+	return m.LinkRefsForItem(item)
+}
+
+// LinkRefsForItem returns references to links related to the requested item.
+func (m *Mapper) LinkRefsForItem(item *pbresource.ID) []*pbresource.Reference {
 	if !resource.EqualType(item.Type, m.itemType) {
 		panic(fmt.Sprintf("expected item type %q got %q",
 			resource.TypeToString(m.itemType),
@@ -157,16 +191,58 @@ func (m *Mapper) LinksForItem(item *pbresource.ID) []*pbresource.Reference {
 	return out
 }
 
+// LinkIDsForItem returns IDs to links related to the requested item.
+func (m *Mapper) LinkIDsForItem(item *pbresource.ID) []*pbresource.ID {
+	if !resource.EqualType(item.Type, m.itemType) {
+		panic(fmt.Sprintf("expected item type %q got %q",
+			resource.TypeToString(m.itemType),
+			resource.TypeToString(item.Type),
+		))
+	}
+
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
+	links, ok := m.itemToLink[resource.NewReferenceKey(item)]
+	if !ok {
+		return nil
+	}
+
+	out := make([]*pbresource.ID, 0, len(links))
+	for l := range links {
+		out = append(out, l.ToID())
+	}
+	return out
+}
+
 // ItemsForLink returns item ids for items related to the provided link.
+// Deprecated: use ItemIDsForLink
 func (m *Mapper) ItemsForLink(link *pbresource.ID) []*pbresource.ID {
+	return m.ItemIDsForLink(link)
+}
+
+// ItemIDsForLink returns item ids for items related to the provided link.
+func (m *Mapper) ItemIDsForLink(link *pbresource.ID) []*pbresource.ID {
 	if !resource.EqualType(link.Type, m.linkType) {
-		panic(fmt.Sprintf("expected type %q got %q",
+		panic(fmt.Sprintf("expected link type %q got %q",
 			resource.TypeToString(m.linkType),
 			resource.TypeToString(link.Type),
 		))
 	}
 
-	return m.itemsByLink(resource.NewReferenceKey(link))
+	return m.itemIDsByLink(resource.NewReferenceKey(link))
+}
+
+// ItemRefsForLink returns item references for items related to the provided link.
+func (m *Mapper) ItemRefsForLink(link *pbresource.ID) []*pbresource.Reference {
+	if !resource.EqualType(link.Type, m.linkType) {
+		panic(fmt.Sprintf("expected link type %q got %q",
+			resource.TypeToString(m.linkType),
+			resource.TypeToString(link.Type),
+		))
+	}
+
+	return m.itemRefsByLink(resource.NewReferenceKey(link))
 }
 
 // MapLink is suitable as a DependencyMapper to map the provided link event to its item.
@@ -180,7 +256,7 @@ func (m *Mapper) MapLink(_ context.Context, _ controller.Runtime, res *pbresourc
 		)
 	}
 
-	itemIDs := m.itemsByLink(resource.NewReferenceKey(link))
+	itemIDs := m.itemIDsByLink(resource.NewReferenceKey(link))
 
 	out := make([]controller.Request, 0, len(itemIDs))
 	for _, item := range itemIDs {
@@ -195,11 +271,8 @@ func (m *Mapper) MapLink(_ context.Context, _ controller.Runtime, res *pbresourc
 	return out, nil
 }
 
-func (m *Mapper) itemsByLink(link resource.ReferenceKey) []*pbresource.ID {
+func (m *Mapper) itemIDsByLink(link resource.ReferenceKey) []*pbresource.ID {
-	m.lock.Lock()
+	items, ok := m.getItemsByLink(link)
-	defer m.lock.Unlock()
-
-	items, ok := m.linkToItem[link]
 	if !ok {
 		return nil
 	}
@@ -210,3 +283,24 @@ func (m *Mapper) itemsByLink(link resource.ReferenceKey) []*pbresource.ID {
 	}
 	return out
 }
+
+func (m *Mapper) itemRefsByLink(link resource.ReferenceKey) []*pbresource.Reference {
+	items, ok := m.getItemsByLink(link)
+	if !ok {
+		return nil
+	}
+
+	out := make([]*pbresource.Reference, 0, len(items))
+	for item := range items {
+		out = append(out, item.ToReference())
+	}
+	return out
+}
+
+func (m *Mapper) getItemsByLink(link resource.ReferenceKey) (map[resource.ReferenceKey]struct{}, bool) {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
+	items, ok := m.linkToItem[link]
+	return items, ok
+}

internal/resource/mappers/bimapper/bimapper_test.go

@@ -49,20 +49,19 @@ func TestMapper(t *testing.T) {
 	wwwRef := newRef(fakeBarType, "www")
 
 	fail1 := rtest.Resource(fakeFooType, "api").Build()
-	fail1_refs := []*pbresource.Reference{
+	fail1Refs := []resource.ReferenceOrID{
 		apiRef,
 		fooRef,
 		barRef,
 	}
 
 	fail2 := rtest.Resource(fakeFooType, "www").Build()
-	fail2_refs := []*pbresource.Reference{
+	fail2Refs := []resource.ReferenceOrID{
 		wwwRef,
 		fooRef,
 	}
 
-	fail1_updated := rtest.Resource(fakeFooType, "api").Build()
+	fail1UpdatedRefs := []resource.ReferenceOrID{
-	fail1_updated_refs := []*pbresource.Reference{
 		apiRef,
 		barRef,
 	}
@@ -101,9 +100,12 @@ func TestMapper(t *testing.T) {
 	requireServicesTracked(t, m, wwwSvc)
 
 	// Actually insert some data.
-	m.TrackItem(fail1.Id, fail1_refs)
+	m.TrackItem(fail1.Id, fail1Refs)
 
-	requireLinksForItem(t, m, fail1.Id, fail1_refs...)
+	// Check links mapping
+	requireLinksForItem(t, m, fail1.Id, fail1Refs...)
+
+	requireLinksForItem(t, m, fail1.Id, fail1Refs...)
 	requireItemsForLink(t, m, apiRef, fail1.Id)
 	requireItemsForLink(t, m, fooRef, fail1.Id)
 	requireItemsForLink(t, m, barRef, fail1.Id)
@@ -116,9 +118,9 @@ func TestMapper(t *testing.T) {
 	requireServicesTracked(t, m, wwwSvc)
 
 	// track it again, no change
-	m.TrackItem(fail1.Id, fail1_refs)
+	m.TrackItem(fail1.Id, fail1Refs)
 
-	requireLinksForItem(t, m, fail1.Id, fail1_refs...)
+	requireLinksForItem(t, m, fail1.Id, fail1Refs...)
 	requireItemsForLink(t, m, apiRef, fail1.Id)
 	requireItemsForLink(t, m, fooRef, fail1.Id)
 	requireItemsForLink(t, m, barRef, fail1.Id)
@@ -131,10 +133,11 @@ func TestMapper(t *testing.T) {
 	requireServicesTracked(t, m, wwwSvc)
 
 	// track new one that overlaps slightly
-	m.TrackItem(fail2.Id, fail2_refs)
+	m.TrackItem(fail2.Id, fail2Refs)
 
-	requireLinksForItem(t, m, fail1.Id, fail1_refs...)
+	// Check links mapping for the new one
-	requireLinksForItem(t, m, fail2.Id, fail2_refs...)
+	requireLinksForItem(t, m, fail1.Id, fail1Refs...)
+	requireLinksForItem(t, m, fail2.Id, fail2Refs...)
 	requireItemsForLink(t, m, apiRef, fail1.Id)
 	requireItemsForLink(t, m, fooRef, fail1.Id, fail2.Id)
 	requireItemsForLink(t, m, barRef, fail1.Id)
@@ -147,10 +150,10 @@ func TestMapper(t *testing.T) {
 	requireServicesTracked(t, m, wwwSvc, fail2.Id)
 
 	// update the original to change it
-	m.TrackItem(fail1_updated.Id, fail1_updated_refs)
+	m.TrackItem(fail1.Id, fail1UpdatedRefs)
 
-	requireLinksForItem(t, m, fail1.Id, fail1_updated_refs...)
+	requireLinksForItem(t, m, fail1.Id, fail1UpdatedRefs...)
-	requireLinksForItem(t, m, fail2.Id, fail2_refs...)
+	requireLinksForItem(t, m, fail2.Id, fail2Refs...)
 	requireItemsForLink(t, m, apiRef, fail1.Id)
 	requireItemsForLink(t, m, fooRef, fail2.Id)
 	requireItemsForLink(t, m, barRef, fail1.Id)
@@ -166,7 +169,7 @@ func TestMapper(t *testing.T) {
 	m.UntrackItem(fail1.Id)
 
 	requireLinksForItem(t, m, fail1.Id)
-	requireLinksForItem(t, m, fail2.Id, fail2_refs...)
+	requireLinksForItem(t, m, fail2.Id, fail2Refs...)
 	requireItemsForLink(t, m, apiRef)
 	requireItemsForLink(t, m, fooRef, fail2.Id)
 	requireItemsForLink(t, m, barRef)
@@ -178,7 +181,16 @@ func TestMapper(t *testing.T) {
 	requireServicesTracked(t, m, barSvc)
 	requireServicesTracked(t, m, wwwSvc, fail2.Id)
 
-	// delete the other one
+	// delete the link
+	m.UntrackLink(newRef(fakeBarType, "www"))
+
+	requireLinksForItem(t, m, fail2.Id, newRef(fakeBarType, "foo"))
+
+	m.UntrackLink(newRef(fakeBarType, "foo"))
+
+	requireLinksForItem(t, m, fail2.Id)
+
+	// delete another item
 	m.UntrackItem(fail2.Id)
 
 	requireLinksForItem(t, m, fail1.Id)
@@ -193,6 +205,88 @@ func TestMapper(t *testing.T) {
 	requireServicesTracked(t, m, fooSvc)
 	requireServicesTracked(t, m, barSvc)
 	requireServicesTracked(t, m, wwwSvc)
+
+	// Reset the mapper and check that its internal maps are empty.
+	m.Reset()
+	require.True(t, m.IsEmpty())
+}
+
+func TestPanics(t *testing.T) {
+	t.Run("new mapper without types", func(t *testing.T) {
+		require.PanicsWithValue(t, "itemType is required", func() {
+			New(nil, nil)
+		})
+
+		require.PanicsWithValue(t, "itemType is required", func() {
+			New(nil, fakeBarType)
+		})
+
+		require.PanicsWithValue(t, "linkType is required", func() {
+			New(fakeFooType, nil)
+		})
+	})
+
+	t.Run("UntrackItem: mismatched type", func(t *testing.T) {
+		m := New(fakeFooType, fakeBarType)
+		require.PanicsWithValue(t, "expected item type \"catalog.v1.Foo\" got \"catalog.v1.Bar\"", func() {
+			// Calling UntrackItem with link type instead of item type
+			m.UntrackItem(rtest.Resource(fakeBarType, "test").ID())
+		})
+	})
+
+	t.Run("TrackItem: mismatched item type", func(t *testing.T) {
+		m := New(fakeFooType, fakeBarType)
+		require.PanicsWithValue(t, "expected item type \"catalog.v1.Foo\" got \"catalog.v1.Bar\"", func() {
+			// Calling UntrackItem with link type instead of item type
+			m.TrackItem(rtest.Resource(fakeBarType, "test").ID(), nil)
+		})
+	})
+
+	t.Run("TrackItem: mismatched link type", func(t *testing.T) {
+		m := New(fakeFooType, fakeBarType)
+		require.PanicsWithValue(t, "expected link type \"catalog.v1.Bar\" got \"catalog.v1.Foo\"", func() {
+			// Calling UntrackItem with link type instead of item type
+			links := []resource.ReferenceOrID{
+				rtest.Resource(fakeFooType, "link").ID(),
+			}
+			m.TrackItem(rtest.Resource(fakeFooType, "test").ID(), links)
+		})
+	})
+
+	t.Run("UntrackLink: mismatched type", func(t *testing.T) {
+		m := New(fakeFooType, fakeBarType)
+		require.PanicsWithValue(t, "expected link type \"catalog.v1.Bar\" got \"catalog.v1.Foo\"", func() {
+			m.UntrackLink(rtest.Resource(fakeFooType, "test").ID())
+		})
+	})
+
+	t.Run("LinkRefsForItem: mismatched type", func(t *testing.T) {
+		m := New(fakeFooType, fakeBarType)
+		require.PanicsWithValue(t, "expected item type \"catalog.v1.Foo\" got \"catalog.v1.Bar\"", func() {
+			m.LinkRefsForItem(rtest.Resource(fakeBarType, "test").ID())
+		})
+	})
+
+	t.Run("LinkRefsForItem: mismatched type", func(t *testing.T) {
+		m := New(fakeFooType, fakeBarType)
+		require.PanicsWithValue(t, "expected item type \"catalog.v1.Foo\" got \"catalog.v1.Bar\"", func() {
+			m.LinkIDsForItem(rtest.Resource(fakeBarType, "test").ID())
+		})
+	})
+
+	t.Run("ItemRefsForLink: mismatched type", func(t *testing.T) {
+		m := New(fakeFooType, fakeBarType)
+		require.PanicsWithValue(t, "expected link type \"catalog.v1.Bar\" got \"catalog.v1.Foo\"", func() {
+			m.ItemRefsForLink(rtest.Resource(fakeFooType, "test").ID())
+		})
+	})
+
+	t.Run("ItemIDsForLink: mismatched type", func(t *testing.T) {
+		m := New(fakeFooType, fakeBarType)
+		require.PanicsWithValue(t, "expected link type \"catalog.v1.Bar\" got \"catalog.v1.Foo\"", func() {
+			m.ItemIDsForLink(rtest.Resource(fakeFooType, "test").ID())
+		})
+	})
 }
 
 func requireServicesTracked(t *testing.T, mapper *Mapper, link *pbresource.Resource, items ...*pbresource.ID) {
@@ -207,39 +301,56 @@ func requireServicesTracked(t *testing.T, mapper *Mapper, link *pbresource.Resou
 
 	require.Len(t, reqs, len(items))
 
+	// Also check items IDs and Refs for link.
+	ids := mapper.ItemIDsForLink(link.Id)
+	require.Len(t, ids, len(items))
+
+	refs := mapper.ItemRefsForLink(link.Id)
+	require.Len(t, refs, len(items))
+
 	for _, item := range items {
 		prototest.AssertContainsElement(t, reqs, controller.Request{ID: item})
+		prototest.AssertContainsElement(t, ids, item)
+		prototest.AssertContainsElement(t, refs, resource.Reference(item, ""))
 	}
 }
 
-func requireLinksForItem(t *testing.T, mapper *Mapper, item *pbresource.ID, links ...*pbresource.Reference) {
-	t.Helper()
-
-	got := mapper.LinksForItem(item)
-
-	prototest.AssertElementsMatch(t, links, got)
-}
-
 func requireItemsForLink(t *testing.T, mapper *Mapper, link *pbresource.Reference, items ...*pbresource.ID) {
 	t.Helper()
 
-	got := mapper.ItemsForLink(resource.IDFromReference(link))
+	got := mapper.ItemIDsForLink(resource.IDFromReference(link))
 
 	prototest.AssertElementsMatch(t, items, got)
 }
 
+func requireLinksForItem(t *testing.T, mapper *Mapper, item *pbresource.ID, links ...resource.ReferenceOrID) {
+	t.Helper()
+
+	var expLinkRefs []*pbresource.Reference
+	var expLinkIDs []*pbresource.ID
+
+	for _, l := range links {
+		expLinkRefs = append(expLinkRefs, &pbresource.Reference{
+			Name:    l.GetName(),
+			Tenancy: l.GetTenancy(),
+			Type:    l.GetType(),
+		})
+		expLinkIDs = append(expLinkIDs, &pbresource.ID{
+			Name:    l.GetName(),
+			Tenancy: l.GetTenancy(),
+			Type:    l.GetType(),
+		})
+	}
+
+	refs := mapper.LinkRefsForItem(item)
+	require.Len(t, refs, len(links))
+	prototest.AssertElementsMatch(t, expLinkRefs, refs)
+
+	ids := mapper.LinkIDsForItem(item)
+	require.Len(t, refs, len(links))
+	prototest.AssertElementsMatch(t, expLinkIDs, ids)
+}
+
 func newRef(typ *pbresource.Type, name string) *pbresource.Reference {
 	return rtest.Resource(typ, name).Reference("")
 }
-
-func newID(typ *pbresource.Type, name string) *pbresource.ID {
-	return rtest.Resource(typ, name).ID()
-}
-
-func defaultTenancy() *pbresource.Tenancy {
-	return &pbresource.Tenancy{
-		Partition: "default",
-		Namespace: "default",
-		PeerName:  "local",
-	}
-}