sets: Add utilities for map of sets (istio#43183)

* add more set operations * add tests * sets: Add utilities for map of sets This introduces two new helper functions to add and remove to map of sets. These are simpler to use and help ensure all callers are correct + efficient.
WillardHu · Feb 8, 2023 · 4974c34 · 4974c34
1 parent c197f97
commit 4974c34
Show file tree

Hide file tree

Showing 10 changed files with 58 additions and 76 deletions.
diff --git a/pilot/pkg/model/push_context.go b/pilot/pkg/model/push_context.go
@@ -728,16 +728,13 @@ func virtualServiceDestinations(v *networking.VirtualService) map[string]sets.Se
 	out := make(map[string]sets.Set[int])
 
 	addDestination := func(host string, port *networking.PortSelector) {
-		if _, ok := out[host]; !ok {
-			out[host] = sets.New[int]()
-		}
+		// Use the value 0 as a sentinel indicating that one of the destinations
+		// in the Virtual Service does not specify a port for this host.
+		pn := 0
 		if port != nil {
-			out[host].Insert(int(port.Number))
-		} else {
-			// Use the value 0 as a sentinel indicating that one of the destinations
-			// in the Virtual Service does not specify a port for this host.
-			out[host].Insert(0)
+			pn = int(port.Number)
 		}
+		sets.InsertOrNew(out, host, pn)
 	}
 
 	for _, h := range v.Http {
@@ -1601,11 +1598,8 @@ func (ps *PushContext) initVirtualServices(env *Environment) error {
 				if gw == constants.IstioMeshGateway {
 					continue
 				}
-				if _, f := ps.virtualServiceIndex.destinationsByGateway[gw]; !f {
-					ps.virtualServiceIndex.destinationsByGateway[gw] = sets.New[string]()
-				}
 				for host := range virtualServiceDestinations(rule) {
-					ps.virtualServiceIndex.destinationsByGateway[gw].Insert(host)
+					sets.InsertOrNew(ps.virtualServiceIndex.destinationsByGateway, gw, host)
 				}
 				addHostsFromMeshConfig(ps, ps.virtualServiceIndex.destinationsByGateway[gw])
 			}

diff --git a/pilot/pkg/model/xds_cache.go b/pilot/pkg/model/xds_cache.go
@@ -189,45 +189,27 @@ func (l *lruCache) onEvict(k string, v cacheValue) {
 
 func (l *lruCache) updateConfigIndex(k string, dependentConfigs []ConfigHash) {
 	for _, cfg := range dependentConfigs {
-		if l.configIndex[cfg] == nil {
-			l.configIndex[cfg] = sets.New[string]()
-		}
-		l.configIndex[cfg].Insert(k)
+		sets.InsertOrNew(l.configIndex, cfg, k)
 	}
 	l.recordDependentConfigSize()
 }
 
 func (l *lruCache) clearConfigIndex(k string, dependentConfigs []ConfigHash) {
 	for _, cfg := range dependentConfigs {
-		index := l.configIndex[cfg]
-		if index != nil {
-			index.Delete(k)
-			if index.IsEmpty() {
-				delete(l.configIndex, cfg)
-			}
-		}
+		sets.DeleteCleanupLast(l.configIndex, cfg, k)
 	}
 	l.recordDependentConfigSize()
 }
 
 func (l *lruCache) updateTypesIndex(k string, dependentTypes []kind.Kind) {
 	for _, t := range dependentTypes {
-		if l.typesIndex[t] == nil {
-			l.typesIndex[t] = sets.New[string]()
-		}
-		l.typesIndex[t].Insert(k)
+		sets.InsertOrNew(l.typesIndex, t, k)
 	}
 }
 
 func (l *lruCache) clearTypesIndex(k string, dependentTypes []kind.Kind) {
 	for _, t := range dependentTypes {
-		index := l.typesIndex[t]
-		if index != nil {
-			index.Delete(k)
-			if index.IsEmpty() {
-				delete(l.typesIndex, t)
-			}
-		}
+		sets.DeleteCleanupLast(l.typesIndex, t, k)
 	}
 }
 

diff --git a/pilot/pkg/networking/core/v1alpha3/cluster.go b/pilot/pkg/networking/core/v1alpha3/cluster.go
@@ -119,10 +119,7 @@ func (configgen *ConfigGeneratorImpl) BuildDeltaClusters(proxy *model.Proxy, upd
 		// WatchedResources.ResourceNames will contain the names of the clusters it is subscribed to. We can
 		// check with the name of our service (cluster names are in the format outbound|<port>||<hostname>).
 		_, _, svcHost, port := model.ParseSubsetKey(cluster)
-		if serviceClusters[string(svcHost)] == nil {
-			serviceClusters[string(svcHost)] = sets.New[string]()
-		}
-		serviceClusters[string(svcHost)].Insert(cluster)
+		sets.InsertOrNew(serviceClusters, string(svcHost), cluster)
 		if servicePorts[string(svcHost)] == nil {
 			servicePorts[string(svcHost)] = make(map[int]string)
 		}

diff --git a/pilot/pkg/networking/grpcgen/cds.go b/pilot/pkg/networking/grpcgen/cds.go
@@ -65,10 +65,7 @@ func newClusterFilter(names []string) map[string]sets.String {
 	for _, name := range names {
 		dir, _, hn, p := model.ParseSubsetKey(name)
 		defaultKey := model.BuildSubsetKey(dir, "", hn, p)
-		if _, ok := filter[defaultKey]; !ok {
-			filter[defaultKey] = sets.New[string]()
-		}
-		filter[defaultKey].Insert(name)
+		sets.InsertOrNew(filter, defaultKey, name)
 	}
 	return filter
 }

diff --git a/pilot/pkg/serviceregistry/kube/controller/pod.go b/pilot/pkg/serviceregistry/kube/controller/pod.go
@@ -262,22 +262,15 @@ func (pc *PodCache) update(ip, key string) {
 func (pc *PodCache) queueEndpointEventOnPodArrival(key, ip string) {
 	pc.Lock()
 	defer pc.Unlock()
-	if _, f := pc.needResync[ip]; !f {
-		pc.needResync[ip] = sets.New(key)
-	} else {
-		pc.needResync[ip].Insert(key)
-	}
+	sets.InsertOrNew(pc.needResync, ip, key)
 	endpointsPendingPodUpdate.Record(float64(len(pc.needResync)))
 }
 
 // endpointDeleted cleans up endpoint from resync endpoint list.
 func (pc *PodCache) endpointDeleted(key string, ip string) {
 	pc.Lock()
 	defer pc.Unlock()
-	delete(pc.needResync[ip], key)
-	if len(pc.needResync[ip]) == 0 {
-		delete(pc.needResync, ip)
-	}
+	sets.DeleteCleanupLast(pc.needResync, ip, key)
 	endpointsPendingPodUpdate.Record(float64(len(pc.needResync)))
 }
 

diff --git a/pilot/pkg/serviceregistry/util/workloadinstances/map.go b/pilot/pkg/serviceregistry/util/workloadinstances/map.go
@@ -24,21 +24,12 @@ type MultiValueMap map[string]sets.String
 
 // Insert adds given (key, value) pair into the map.
 func (m MultiValueMap) Insert(key, value string) MultiValueMap {
-	if values, exists := m[key]; exists {
-		values.Insert(value)
-		return m
-	}
-	m[key] = sets.New(value)
+	sets.InsertOrNew(m, key, value)
 	return m
 }
 
 // Delete removes given (key, value) pair out of the map.
 func (m MultiValueMap) Delete(key, value string) MultiValueMap {
-	if values, exists := m[key]; exists {
-		values.Delete(value)
-		if values.IsEmpty() {
-			delete(m, key)
-		}
-	}
+	sets.DeleteCleanupLast(m, key, value)
 	return m
 }
diff --git a/pilot/pkg/status/distribution/reporter.go b/pilot/pkg/status/distribution/reporter.go
@@ -316,22 +316,14 @@ func (r *Reporter) processEvent(conID string, distributionType xds.EventType, no
 	}
 	// touch
 	r.status[key] = version
-	if _, ok := r.reverseStatus[version]; !ok {
-		r.reverseStatus[version] = sets.New[string]()
-	}
-	r.reverseStatus[version].Insert(key)
+	sets.InsertOrNew(r.reverseStatus, version, key)
 }
 
 // This is a helper function for keeping our reverseStatus map in step with status.
 // must have write lock before calling.
 func (r *Reporter) deleteKeyFromReverseMap(key string) {
 	if old, ok := r.status[key]; ok {
-		if keys, ok := r.reverseStatus[old]; ok {
-			keys.Delete(key)
-			if keys.IsEmpty() {
-				delete(r.reverseStatus, old)
-			}
-		}
+		sets.DeleteCleanupLast(r.reverseStatus, old, key)
 	}
 }
 

diff --git a/pkg/util/sets/set.go b/pkg/util/sets/set.go
@@ -216,3 +216,28 @@ func (s Set[T]) Len() int {
 func (s Set[T]) IsEmpty() bool {
 	return len(s) == 0
 }
+
+// InsertOrNew inserts t into the set if the set exists, or returns a new set with t if not.
+// Works well with DeleteCleanupLast.
+// Example:
+//
+//	InsertOrNew(m, key, value)
+func InsertOrNew[K comparable, T comparable](m map[K]Set[T], k K, v T) {
+	s, f := m[k]
+	if !f {
+		m[k] = New(v)
+	} else {
+		s.Insert(v)
+	}
+}
+
+// DeleteCleanupLast removes an element from a set in a map of sets, deleting the key from the map if there are no keys left.
+// Works well with InsertOrNew.
+// Example:
+//
+//	sets.DeleteCleanupLast(m, key, value)
+func DeleteCleanupLast[K comparable, T comparable](m map[K]Set[T], k K, v T) {
+	if m[k].Delete(v).IsEmpty() {
+		delete(m, k)
+	}
+}
diff --git a/pkg/util/sets/set_test.go b/pkg/util/sets/set_test.go
@@ -324,3 +324,17 @@ func BenchmarkSet(b *testing.B) {
 		}
 	})
 }
+
+func TestMapOfSet(t *testing.T) {
+	m := map[int]String{}
+	InsertOrNew(m, 1, "a")
+	InsertOrNew(m, 1, "b")
+	InsertOrNew(m, 2, "c")
+	assert.Equal(t, m, map[int]String{1: New("a", "b"), 2: New("c")})
+
+	DeleteCleanupLast(m, 1, "a")
+	assert.Equal(t, m, map[int]String{1: New("b"), 2: New("c")})
+	DeleteCleanupLast(m, 1, "b")
+	DeleteCleanupLast(m, 1, "not found")
+	assert.Equal(t, m, map[int]String{2: New("c")})
+}
diff --git a/tools/bug-report/pkg/bugreport/bugreport.go b/tools/bug-report/pkg/bugreport/bugreport.go
@@ -237,10 +237,7 @@ func getIstioVersions(kubeconfig, configContext, istioNamespace string, revision
 			continue
 		}
 		for _, pi := range *proxyInfo {
-			if proxyVersionsMap[revision] == nil {
-				proxyVersionsMap[revision] = sets.New[string]()
-			}
-			proxyVersionsMap[revision].Insert(pi.IstioVersion)
+			sets.InsertOrNew(proxyVersionsMap, revision, pi.IstioVersion)
 		}
 	}
 	for revision, vmap := range proxyVersionsMap {