fix(zrpc): aggregate kube resolver addresses by EndpointSlice

zrpc/resolver/internal/kube/eventhandler.go

@@ -13,22 +13,22 @@ var _ cache.ResourceEventHandler = (*EventHandler)(nil)
 
 // EventHandler is ResourceEventHandler implementation.
 type EventHandler struct {
-	update    func([]string)
-	endpoints map[string]lang.PlaceholderType
-	lock      sync.Mutex
+	update         func([]string)
+	endpointSlices map[string]map[string]lang.PlaceholderType
+	lock           sync.Mutex
 }
 
 // NewEventHandler returns an EventHandler.
 func NewEventHandler(update func([]string)) *EventHandler {
 	return &EventHandler{
-		update:    update,
-		endpoints: make(map[string]lang.PlaceholderType),
+		update:         update,
+		endpointSlices: make(map[string]map[string]lang.PlaceholderType),
 	}
 }
 
 // OnAdd handles the endpoints add events.
 func (h *EventHandler) OnAdd(obj any, _ bool) {
-	endpoints, ok := obj.(*v1.EndpointSlice)
+	endpoints, ok := parseEndpointSlice(obj)
 	if !ok {
 		logx.Errorf("%v is not an object with type *v1.EndpointSlice", obj)
 		return
@@ -37,24 +37,16 @@ func (h *EventHandler) OnAdd(obj any, _ bool) {
 	h.lock.Lock()
 	defer h.lock.Unlock()
 
-	var changed bool
-	for _, point := range endpoints.Endpoints {
-		for _, address := range point.Addresses {
-			if _, ok := h.endpoints[address]; !ok {
-				h.endpoints[address] = lang.Placeholder
-				changed = true
-			}
-		}
-	}
-
-	if changed {
+	old := h.endpoints()
+	h.updateEndpointSlice(endpoints)
+	if diff(old, h.endpoints()) {
 		h.notify()
 	}
 }
 
 // OnDelete handles the endpoints delete events.
 func (h *EventHandler) OnDelete(obj any) {
-	endpoints, ok := obj.(*v1.EndpointSlice)
+	endpoints, ok := parseEndpointSlice(obj)
 	if !ok {
 		logx.Errorf("%v is not an object with type *v1.EndpointSlice", obj)
 		return
@@ -63,30 +55,27 @@ func (h *EventHandler) OnDelete(obj any) {
 	h.lock.Lock()
 	defer h.lock.Unlock()
 
-	var changed bool
-	for _, point := range endpoints.Endpoints {
-		for _, address := range point.Addresses {
-			if _, ok := h.endpoints[address]; ok {
-				delete(h.endpoints, address)
-				changed = true
-			}
-		}
+	old := h.endpoints()
+	if key := endpointSliceKey(endpoints); len(key) > 0 {
+		delete(h.endpointSlices, key)
+	} else {
+		h.deleteEndpointAddresses(endpoints)
 	}
 
-	if changed {
+	if diff(old, h.endpoints()) {
 		h.notify()
 	}
 }
 
 // OnUpdate handles the endpoints update events.
 func (h *EventHandler) OnUpdate(oldObj, newObj any) {
-	oldEndpointSlices, ok := oldObj.(*v1.EndpointSlice)
+	oldEndpointSlices, ok := parseEndpointSlice(oldObj)
 	if !ok {
 		logx.Errorf("%v is not an object with type *v1.EndpointSlice", oldObj)
 		return
 	}
 
-	newEndpointSlices, ok := newObj.(*v1.EndpointSlice)
+	newEndpointSlices, ok := parseEndpointSlice(newObj)
 	if !ok {
 		logx.Errorf("%v is not an object with type *v1.EndpointSlice", newObj)
 		return
@@ -104,29 +93,82 @@ func (h *EventHandler) Update(endpoints *v1.EndpointSlice) {
 	h.lock.Lock()
 	defer h.lock.Unlock()
 
-	old := h.endpoints
-	h.endpoints = make(map[string]lang.PlaceholderType)
+	old := h.endpoints()
+	h.updateEndpointSlice(endpoints)
+	if diff(old, h.endpoints()) {
+		h.notify()
+	}
+}
+
+func (h *EventHandler) updateEndpointSlice(endpoints *v1.EndpointSlice) {
+	h.endpointSlices[endpointSliceKey(endpoints)] = endpointAddresses(endpoints)
+}
+
+func (h *EventHandler) deleteEndpointAddresses(endpoints *v1.EndpointSlice) {
 	for _, point := range endpoints.Endpoints {
 		for _, address := range point.Addresses {
-			h.endpoints[address] = lang.Placeholder
+			for key, addresses := range h.endpointSlices {
+				delete(addresses, address)
+				if len(addresses) == 0 {
+					delete(h.endpointSlices, key)
+				}
+			}
 		}
 	}
+}
 
-	if diff(old, h.endpoints) {
-		h.notify()
+func (h *EventHandler) endpoints() map[string]lang.PlaceholderType {
+	endpoints := make(map[string]lang.PlaceholderType)
+	for _, slice := range h.endpointSlices {
+		for address := range slice {
+			endpoints[address] = lang.Placeholder
+		}
 	}
+
+	return endpoints
 }
 
 func (h *EventHandler) notify() {
-	targets := make([]string, 0, len(h.endpoints))
+	endpoints := h.endpoints()
+	targets := make([]string, 0, len(endpoints))
 
-	for k := range h.endpoints {
+	for k := range endpoints {
 		targets = append(targets, k)
 	}
 
 	h.update(targets)
 }
 
+func parseEndpointSlice(obj any) (*v1.EndpointSlice, bool) {
+	switch endpoints := obj.(type) {
+	case *v1.EndpointSlice:
+		return endpoints, true
+	case cache.DeletedFinalStateUnknown:
+		return parseEndpointSlice(endpoints.Obj)
+	default:
+		return nil, false
+	}
+}
+
+func endpointAddresses(endpoints *v1.EndpointSlice) map[string]lang.PlaceholderType {
+	addresses := make(map[string]lang.PlaceholderType)
+	for _, point := range endpoints.Endpoints {
+		for _, address := range point.Addresses {
+			addresses[address] = lang.Placeholder
+		}
+	}
+
+	return addresses
+}
+
+func endpointSliceKey(endpoints *v1.EndpointSlice) string {
+	if len(endpoints.Namespace) > 0 || len(endpoints.Name) > 0 {
+		return endpoints.Namespace + "/" + endpoints.Name
+	}
+
+	return string(endpoints.UID)
+}
+
 func diff(o, n map[string]lang.PlaceholderType) bool {
 	if len(o) != len(n) {
 		return true

zrpc/resolver/internal/kube/eventhandler_test.go

@@ -6,6 +6,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	discoveryv1 "k8s.io/api/discovery/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/tools/cache"
 )
 
 func TestAdd(t *testing.T) {
@@ -62,6 +63,34 @@ func TestDelete(t *testing.T) {
 	assert.ElementsMatch(t, []string{"0.0.0.3"}, endpoints)
 }
 
+func TestDeleteNamedEndpointSlice(t *testing.T) {
+	var endpoints []string
+	h := NewEventHandler(func(change []string) {
+		endpoints = change
+	})
+	h.OnAdd(newEndpointSlice("slice-a", "1", "0.0.0.1", "0.0.0.2"), false)
+	h.OnAdd(newEndpointSlice("slice-b", "1", "0.0.0.3"), false)
+
+	h.OnDelete(newEndpointSlice("slice-a", "1", "0.0.0.1", "0.0.0.2"))
+
+	assert.ElementsMatch(t, []string{"0.0.0.3"}, endpoints)
+}
+
+func TestDeleteEndpointSliceTombstone(t *testing.T) {
+	var endpoints []string
+	h := NewEventHandler(func(change []string) {
+		endpoints = change
+	})
+	h.OnAdd(newEndpointSlice("slice-a", "1", "0.0.0.1"), false)
+	h.OnAdd(newEndpointSlice("slice-b", "1", "0.0.0.2"), false)
+
+	h.OnDelete(cache.DeletedFinalStateUnknown{
+		Obj: newEndpointSlice("slice-a", "1", "0.0.0.1"),
+	})
+
+	assert.ElementsMatch(t, []string{"0.0.0.2"}, endpoints)
+}
+
 func TestUpdate(t *testing.T) {
 	var endpoints []string
 	h := NewEventHandler(func(change []string) {
@@ -228,3 +257,47 @@ func TestUpdateNoChangeWithDifferentVersion(t *testing.T) {
 	})
 	assert.ElementsMatch(t, []string{"0.0.0.1", "0.0.0.2"}, endpoints)
 }
+
+func TestUpdateEmptyNamedEndpointSliceKeepsOtherSlices(t *testing.T) {
+	var endpoints []string
+	h := NewEventHandler(func(change []string) {
+		endpoints = change
+	})
+	h.OnAdd(newEndpointSlice("slice-a", "1", "0.0.0.1"), false)
+	h.OnAdd(newEndpointSlice("slice-b", "1", "0.0.0.2"), false)
+
+	h.OnUpdate(newEndpointSlice("slice-a", "1", "0.0.0.1"), newEndpointSlice("slice-a", "2"))
+
+	assert.ElementsMatch(t, []string{"0.0.0.2"}, endpoints)
+}
+
+func TestUpdateNamedEndpointSliceAggregatesAllSlices(t *testing.T) {
+	var endpoints []string
+	h := NewEventHandler(func(change []string) {
+		endpoints = change
+	})
+	h.OnAdd(newEndpointSlice("slice-a", "1", "0.0.0.1"), false)
+	h.OnAdd(newEndpointSlice("slice-b", "1", "0.0.0.2"), false)
+
+	h.OnUpdate(newEndpointSlice("slice-a", "1", "0.0.0.1"), newEndpointSlice("slice-a", "2", "0.0.0.3"))
+
+	assert.ElementsMatch(t, []string{"0.0.0.2", "0.0.0.3"}, endpoints)
+}
+
+func newEndpointSlice(name, version string, addresses ...string) *discoveryv1.EndpointSlice {
+	endpoints := make([]discoveryv1.Endpoint, 0, len(addresses))
+	for _, address := range addresses {
+		endpoints = append(endpoints, discoveryv1.Endpoint{
+			Addresses: []string{address},
+		})
+	}
+
+	return &discoveryv1.EndpointSlice{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:            name,
+			Namespace:       "default",
+			ResourceVersion: version,
+		},
+		Endpoints: endpoints,
+	}
+}