krt: Add JoinWithMergeCollection (#58332)

* Add JoinWithMergeCollection

* Check if Equal with values

Use values when comparing for equality so either form of Equals can be
used.

* code review

- wait for collections to sync before registering
- remove dupe metadata assignment
- remove incorrect comments

Author: fraenkel

pkg/kube/krt/mergejoin.go

@@ -18,17 +18,24 @@ import (
 	"fmt"
 	"sync"
 
+	"k8s.io/client-go/tools/cache"
+
+	"istio.io/istio/pkg/kube"
 	"istio.io/istio/pkg/kube/controllers"
 	istiolog "istio.io/istio/pkg/log"
 	"istio.io/istio/pkg/maps"
+	"istio.io/istio/pkg/ptr"
 	"istio.io/istio/pkg/queue"
+	"istio.io/istio/pkg/slices"
 	"istio.io/istio/pkg/util/sets"
 )
 
 type collectionLister[T any] interface {
 	getCollections() []Collection[T]
 }
 
+var _ internalCollection[any] = &mergejoin[any]{}
+
 type mergejoin[T any] struct {
 	id             collectionUID
 	collectionName string
@@ -169,6 +176,42 @@ func (j *mergejoin[T]) GetKey(k string) *T {
 	return nil
 }
 
+func (j *mergejoin[T]) Register(f func(e Event[T])) HandlerRegistration {
+	return registerHandlerAsBatched(j, f)
+}
+
+func (j *mergejoin[T]) RegisterBatch(f func(e []Event[T]), runExistingState bool) HandlerRegistration {
+	j.mu.Lock()
+	defer j.mu.Unlock()
+	if !runExistingState {
+		// If we don't to run the initial state this is simple, we just register the handler.
+		return j.eventHandlers.Insert(f, j, nil, j.stop)
+	}
+
+	// We need to run the initial state, but we don't want to get duplicate events.
+	// We should get "ADD initialObject1, ADD initialObjectN, UPDATE someLaterUpdate" without mixing the initial ADDs
+	// Create ADDs for the current state of the merge cache
+	events := make([]Event[T], 0, len(j.outputs))
+	for _, o := range j.outputs {
+		events = append(events, Event[T]{
+			New:   &o,
+			Event: controllers.EventAdd,
+		})
+	}
+
+	// Send out all the initial objects to the handler. We will then unlock the new events so it gets the future updates.
+	return j.eventHandlers.Insert(f, j, events, j.stop)
+}
+
+// nolint: unused // (not true, its to implement an interface)
+func (j *mergejoin[T]) dump() CollectionDump {
+	return CollectionDump{
+		Outputs: eraseMap(slices.GroupUnique(j.List(), getTypedKey)),
+		Synced:  j.HasSynced(),
+		Inputs:  nil,
+	}
+}
+
 func (j *mergejoin[T]) onSubCollectionEventHandler(o []Event[T]) {
 	var events []Event[T]
 	j.mu.Lock()
@@ -205,7 +248,7 @@ func (j *mergejoin[T]) onSubCollectionEventHandler(o []Event[T]) {
 			oldObj := ev.Old
 			if oldObj != nil && newObj != nil {
 				// Update event
-				if Equal(ev.New, ev.Old) {
+				if Equal(*ev.New, *ev.Old) {
 					// NOP change, skip
 					continue
 				}
@@ -313,3 +356,80 @@ func (j *mergejoin[T]) updateIndexLocked(e Event[T], key Key[T]) {
 		}
 	}
 }
+
+func (j *mergejoin[T]) runQueue() {
+	// Wait until all underlying collections are synced before registering
+	syncers := slices.Map(j.collections.getCollections(), func(c Collection[T]) cache.InformerSynced {
+		return c.HasSynced
+	})
+	if !kube.WaitForCacheSync(j.collectionName, j.stop, syncers...) {
+		return
+	}
+
+	// Register with the list of collections
+	regs := []HandlerRegistration{}
+	for _, c := range j.collections.getCollections() {
+		regs = append(regs, c.RegisterBatch(func(events []Event[T]) {
+			j.queue.Push(func() error {
+				j.onSubCollectionEventHandler(events)
+				return nil
+			})
+		}, true))
+	}
+
+	syncers = slices.Map(regs, func(r HandlerRegistration) cache.InformerSynced {
+		return r.HasSynced
+	})
+	if !kube.WaitForCacheSync(j.collectionName, j.stop, syncers...) {
+		return
+	}
+
+	j.queue.Run(j.stop)
+}
+
+type collections[T any] []Collection[T]
+
+// nolint: unused // (not true, its to implement an interface)
+func (jc collections[T]) getCollections() []Collection[T] {
+	return jc
+}
+
+func JoinWithMergeCollection[T any](cs []Collection[T], merge func(ts []T) *T, opts ...CollectionOption) Collection[T] {
+	o := buildCollectionOptions(opts...)
+	if o.name == "" {
+		o.name = fmt.Sprintf("JoinWithMerge[%v]", ptr.TypeName[T]())
+	}
+
+	synced := make(chan struct{})
+
+	j := &mergejoin[T]{
+		collectionName: o.name,
+		collections:    collections[T](cs),
+		id:             nextUID(),
+		log:            log.WithLabels("owner", o.name),
+		outputs:        make(map[Key[T]]T),
+		indexes:        make(map[string]joinCollectionIndex[T]),
+		eventHandlers:  newHandlerSet[T](),
+		metadata:       o.metadata,
+		merge:          merge,
+		synced:         synced,
+		stop:           o.stop,
+		syncer: channelSyncer{
+			name:   o.name,
+			synced: synced,
+		},
+	}
+
+	maybeRegisterCollectionForDebugging(j, o.debugger)
+
+	// Create our queue. When it syncs (that is, all items that were present when Run() was called), we mark ourselves as synced.
+	j.queue = queue.NewWithSync(func() {
+		close(j.synced)
+		j.log.Infof("%v synced (uid %v)", j.name(), j.uid())
+	}, j.collectionName)
+
+	// The queue will process the initial state and mark ourselves as synced (from the NewWithSync callback)
+	go j.runQueue()
+
+	return j
+}