cleanup and tests

tempodb/blocklist/list.go

@@ -4,7 +4,6 @@ import (
 	"slices"
 	"sync"
 
-	"github.com/google/uuid"
 	"github.com/grafana/tempo/tempodb/backend"
 )
 
@@ -126,78 +125,70 @@ func (l *List) Update(tenantID string, add []*backend.BlockMeta, remove []*backe
 // updateInternal exists to do the work of applying updates to held PerTenant and PerTenantCompacted maps
 // it must be called under lock
 func (l *List) updateInternal(tenantID string, add []*backend.BlockMeta, remove []*backend.BlockMeta, compactedAdd []*backend.CompactedBlockMeta, compactedRemove []*backend.CompactedBlockMeta) {
-	// ******** Regular blocks ********
-	blocklist := l.metas[tenantID]
-
-	matchedRemovals := make(map[uuid.UUID]struct{})
-	for _, b := range blocklist {
-		for _, rem := range remove {
-			if b.BlockID == rem.BlockID {
-				matchedRemovals[(uuid.UUID)(rem.BlockID)] = struct{}{}
-				break
-			}
+	hasID := func(id backend.UUID) func(*backend.BlockMeta) bool {
+		return func(b *backend.BlockMeta) bool {
+			return b.BlockID == id
 		}
 	}
 
-	existingMetas := make(map[uuid.UUID]struct{})
-	newblocklist := make([]*backend.BlockMeta, 0, len(blocklist)-len(matchedRemovals)+len(add))
-	// rebuild the blocklist dropping all removals
-	for _, b := range blocklist {
-		existingMetas[(uuid.UUID)(b.BlockID)] = struct{}{}
-		if _, ok := matchedRemovals[(uuid.UUID)(b.BlockID)]; !ok {
-			newblocklist = append(newblocklist, b)
+	hasIDC := func(id backend.UUID) func(*backend.CompactedBlockMeta) bool {
+		return func(b *backend.CompactedBlockMeta) bool {
+			return b.BlockID == id
 		}
 	}
-	// add new blocks (only if they don't already exist) and weren't also removed
-	for _, b := range add {
-		if _, ok := existingMetas[(uuid.UUID)(b.BlockID)]; ok {
-			continue
+
+	// ******** Regular blocks ********
+	if len(add) > 0 || len(remove) > 0 || len(compactedAdd) > 0 || len(compactedRemove) > 0 {
+		var (
+			existing = l.metas[tenantID]
+			final    = make([]*backend.BlockMeta, 0, max(0, len(existing)+len(add)-len(remove)))
+		)
+
+		// rebuild dropping all removals
+		for _, b := range existing {
+			if slices.ContainsFunc(remove, hasID(b.BlockID)) {
+				continue
+			}
+			final = append(final, b)
 		}
+		// add new if they don't already exist and weren't also removed
+		for _, b := range add {
+			if slices.ContainsFunc(final, hasID(b.BlockID)) ||
+				slices.ContainsFunc(remove, hasID(b.BlockID)) ||
+				slices.ContainsFunc(compactedAdd, hasIDC(b.BlockID)) ||
+				slices.ContainsFunc(compactedRemove, hasIDC(b.BlockID)) {
+				continue
+			}
 
-		if slices.ContainsFunc(remove, func(b2 *backend.BlockMeta) bool {
-			return b.BlockID == b2.BlockID
-		}) {
-			continue
+			final = append(final, b)
 		}
-		newblocklist = append(newblocklist, b)
-	}
 
-	l.metas[tenantID] = newblocklist
+		l.metas[tenantID] = final
+	}
 
 	// ******** Compacted blocks ********
-	compactedBlocklist := l.compactedMetas[tenantID]
-
-	compactedRemovals := map[uuid.UUID]struct{}{}
-	for _, c := range compactedBlocklist {
-		for _, rem := range compactedRemove {
-			if c.BlockID == rem.BlockID {
-				compactedRemovals[(uuid.UUID)(rem.BlockID)] = struct{}{}
-				break
+	if len(compactedAdd) > 0 || len(compactedRemove) > 0 {
+		var (
+			existing = l.compactedMetas[tenantID]
+			final    = make([]*backend.CompactedBlockMeta, 0, max(0, len(existing)+len(compactedAdd)-len(compactedRemove)))
+		)
+
+		// rebuild dropping all removals
+		for _, b := range existing {
+			if slices.ContainsFunc(compactedRemove, hasIDC(b.BlockID)) {
+				continue
 			}
+			final = append(final, b)
 		}
-	}
-
-	existingMetas = make(map[uuid.UUID]struct{})
-	newCompactedBlocklist := make([]*backend.CompactedBlockMeta, 0, len(compactedBlocklist)-len(compactedRemovals)+len(compactedAdd))
-	// rebuild the blocklist dropping all removals
-	for _, b := range compactedBlocklist {
-		existingMetas[(uuid.UUID)(b.BlockID)] = struct{}{}
-		if _, ok := compactedRemovals[(uuid.UUID)(b.BlockID)]; !ok {
-			newCompactedBlocklist = append(newCompactedBlocklist, b)
-		}
-	}
-	for _, b := range compactedAdd {
-		if _, ok := existingMetas[(uuid.UUID)(b.BlockID)]; ok {
-			continue
+		// add new if they don't already exist and weren't also removed
+		for _, b := range compactedAdd {
+			if slices.ContainsFunc(existing, hasIDC(b.BlockID)) ||
+				slices.ContainsFunc(compactedRemove, hasIDC(b.BlockID)) {
+				continue
+			}
+			final = append(final, b)
 		}
 
-		if slices.ContainsFunc(compactedRemove, func(b2 *backend.CompactedBlockMeta) bool {
-			return b.BlockID == b2.BlockID
-		}) {
-			continue
-		}
-		newCompactedBlocklist = append(newCompactedBlocklist, b)
+		l.compactedMetas[tenantID] = final
 	}
-
-	l.compactedMetas[tenantID] = newCompactedBlocklist
 }