Optimize AggregatePrefixes performance

pkg/provider/loadbalancer/iputil/internal/bits.go

@@ -0,0 +1,37 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package internal
+
+// setBitAt sets the bit at the i-th position in the byte slice to the given value.
+// Panics if the index is out of bounds.
+// For example,
+// - setBitAt([0x00, 0x00], 8, 1) returns [0x00, 0b1000_0000].
+// - setBitAt([0xff, 0xff], 0, 0) returns [0b0111_1111, 0xff].
+func setBitAt(bytes []byte, i int, bit uint8) {
+	if bit == 1 {
+		bytes[i/8] |= 1 << (7 - i%8)
+	} else {
+		bytes[i/8] &^= 1 << (7 - i%8)
+	}
+}
+
+// bitAt returns the bit at the i-th position in the byte slice.
+// The return value is either 0 or 1 as uint8.
+// Panics if the index is out of bounds.
+func bitAt(bytes []byte, i int) uint8 {
+	return bytes[i/8] >> (7 - i%8) & 1
+}

pkg/provider/loadbalancer/iputil/internal/bits_test.go

@@ -0,0 +1,103 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package internal
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_bitAt(t *testing.T) {
+	bytes := []byte{0b1010_1010, 0b0101_0101}
+	assert.Equal(t, uint8(1), bitAt(bytes, 0))
+	assert.Equal(t, uint8(0), bitAt(bytes, 1))
+	assert.Equal(t, uint8(1), bitAt(bytes, 2))
+	assert.Equal(t, uint8(0), bitAt(bytes, 3))
+
+	assert.Equal(t, uint8(1), bitAt(bytes, 4))
+	assert.Equal(t, uint8(0), bitAt(bytes, 5))
+	assert.Equal(t, uint8(1), bitAt(bytes, 6))
+	assert.Equal(t, uint8(0), bitAt(bytes, 7))
+
+	assert.Equal(t, uint8(0), bitAt(bytes, 8))
+	assert.Equal(t, uint8(1), bitAt(bytes, 9))
+	assert.Equal(t, uint8(0), bitAt(bytes, 10))
+	assert.Equal(t, uint8(1), bitAt(bytes, 11))
+
+	assert.Equal(t, uint8(0), bitAt(bytes, 12))
+	assert.Equal(t, uint8(1), bitAt(bytes, 13))
+	assert.Equal(t, uint8(0), bitAt(bytes, 14))
+	assert.Equal(t, uint8(1), bitAt(bytes, 15))
+
+	assert.Panics(t, func() { bitAt(bytes, 16) })
+}
+
+func Test_setBitAt(t *testing.T) {
+	tests := []struct {
+		name     string
+		initial  []byte
+		index    int
+		bit      uint8
+		expected []byte
+	}{
+		{
+			name:     "Set first bit to 1",
+			initial:  []byte{0b0000_0000},
+			index:    0,
+			bit:      1,
+			expected: []byte{0b1000_0000},
+		},
+		{
+			name:     "Set last bit to 1",
+			initial:  []byte{0b0000_0000},
+			index:    7,
+			bit:      1,
+			expected: []byte{0b0000_0001},
+		},
+		{
+			name:     "Set middle bit to 1",
+			initial:  []byte{0b0000_0000},
+			index:    4,
+			bit:      1,
+			expected: []byte{0b0000_1000},
+		},
+		{
+			name:     "Set bit to 0",
+			initial:  []byte{0b1111_1111},
+			index:    3,
+			bit:      0,
+			expected: []byte{0b1110_1111},
+		},
+		{
+			name:     "Set bit in second byte",
+			initial:  []byte{0b0000_0000, 0b0000_0000},
+			index:    9,
+			bit:      1,
+			expected: []byte{0b0000_0000, 0b0100_0000},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			setBitAt(tt.initial, tt.index, tt.bit)
+			assert.Equal(t, tt.expected, tt.initial)
+		})
+	}
+
+	assert.Panics(t, func() { setBitAt([]byte{0x00}, 8, 1) })
+}

pkg/provider/loadbalancer/iputil/internal/prefix.go

@@ -18,6 +18,7 @@ package internal
 
 import (
 	"net/netip"
+	"sort"
 )
 
 // ListAddresses returns all IP addresses contained within the given prefixes.
@@ -33,3 +34,84 @@ func ListAddresses(prefixes ...netip.Prefix) []netip.Addr {
 	}
 	return rv
 }
+
+// ContainsPrefix checks if prefix p fully contains prefix o.
+// It returns true if o is a subset of p, meaning all addresses in o are also in p.
+// This is true when p overlaps with o and p has fewer or equal number of bits than o.
+func ContainsPrefix(p netip.Prefix, o netip.Prefix) bool {
+	return p.Overlaps(o) && p.Bits() <= o.Bits()
+}
+
+// IsAdjacentPrefixes checks if two prefixes are adjacent and can be merged into a single prefix.
+// Two prefixes are considered adjacent if they have the same length and their
+// addresses are consecutive.
+//
+// Examples:
+//   - Adjacent:     192.168.0.0/32 and 192.168.0.1/32
+//   - Adjacent:     192.168.0.0/24 and 192.168.1.0/24
+//   - Not adjacent: 192.168.0.1/32 and 192.168.0.2/32 (cannot merge)
+//   - Not adjacent: 192.168.0.0/24 and 192.168.0.0/25 (different lengths)
+func IsAdjacentPrefixes(p1, p2 netip.Prefix) bool {
+	if p1.Bits() != p2.Bits() {
+		return false
+	}
+
+	p1Bytes := p1.Addr().AsSlice()
+	p2Bytes := p2.Addr().AsSlice()
+
+	if bitAt(p1Bytes, p1.Bits()-1) == 0 {
+		setBitAt(p1Bytes, p1.Bits()-1, 1)
+		addr, _ := netip.AddrFromSlice(p1Bytes)
+		return addr == p2.Addr()
+	} else {
+		setBitAt(p2Bytes, p2.Bits()-1, 1)
+		addr, _ := netip.AddrFromSlice(p2Bytes)
+		return addr == p1.Addr()
+	}
+}
+
+// AggregatePrefixesForSingleIPFamily merges overlapping or adjacent prefixes into a single prefix.
+// The input prefixes must be the same IP family (IPv4 or IPv6).
+// For example,
+// - [192.168.0.0/32, 192.168.0.1/32] -> [192.168.0.0/31] (adjacent)
+// - [192.168.0.0/24, 192.168.0.1/32] -> [192.168.1.0/24] (overlapping)
+func AggregatePrefixesForSingleIPFamily(prefixes []netip.Prefix) []netip.Prefix {
+	if len(prefixes) <= 1 {
+		return prefixes
+	}
+
+	sort.Slice(prefixes, func(i, j int) bool {
+		if prefixes[i].Addr() == prefixes[j].Addr() {
+			return prefixes[i].Bits() < prefixes[j].Bits()
+		}
+		return prefixes[i].Addr().Less(prefixes[j].Addr())
+	})
+
+	var rv = []netip.Prefix{
+		prefixes[0],
+	}
+
+	for i := 1; i < len(prefixes); {
+		last, p := rv[len(rv)-1], prefixes[i]
+		if ContainsPrefix(last, p) {
+			// Skip overlapping prefixes
+			i++
+			continue
+		}
+		rv = append(rv, p)
+
+		// Merge adjacent prefixes if possible
+		for len(rv) >= 2 {
+			p1, p2 := rv[len(rv)-2], rv[len(rv)-1]
+			if !IsAdjacentPrefixes(p1, p2) {
+				break
+			}
+
+			bits := p1.Bits() - 1
+			p, _ := p1.Addr().Prefix(bits)
+			rv = rv[:len(rv)-2]
+			rv = append(rv, p)
+		}
+	}
+	return rv
+}

pkg/provider/loadbalancer/iputil/internal/prefix_test.go

@@ -17,6 +17,8 @@ limitations under the License.
 package internal
 
 import (
+	"fmt"
+	"math/rand"
 	"net/netip"
 	"testing"
 
@@ -84,3 +86,49 @@ func TestListAddresses(t *testing.T) {
 		})
 	}
 }
+
+func benchmarkPrefixFixtures() []netip.Prefix {
+	var rv []netip.Prefix
+	for i := 0; i <= 255; i++ {
+		for j := 0; j <= 255; j++ {
+			rv = append(rv, netip.MustParsePrefix(fmt.Sprintf("192.168.%d.%d/32", i, j)))
+		}
+	}
+	rand.Shuffle(len(rv), func(i, j int) {
+		rv[i], rv[j] = rv[j], rv[i]
+	})
+
+	return rv
+}
+
+func BenchmarkAggregatePrefixesDefault(b *testing.B) {
+	prefixes := benchmarkPrefixFixtures()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		actual := AggregatePrefixesForSingleIPFamily(prefixes)
+		assert.Len(b, actual, 1)
+		assert.Equal(b, []netip.Prefix{
+			netip.MustParsePrefix("192.168.0.0/16"),
+		}, actual)
+	}
+}
+
+func BenchmarkAggregatePrefixesUsingPrefixTree(b *testing.B) {
+	do := func(prefixes []netip.Prefix) []netip.Prefix {
+		tree := NewPrefixTreeForIPv4()
+		for _, p := range prefixes {
+			tree.Add(p)
+		}
+		return tree.List()
+	}
+
+	prefixes := benchmarkPrefixFixtures()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		actual := do(prefixes)
+		assert.Len(b, actual, 1)
+		assert.Equal(b, []netip.Prefix{
+			netip.MustParsePrefix("192.168.0.0/16"),
+		}, actual)
+	}
+}

pkg/provider/loadbalancer/iputil/internal/prefixtree.go

@@ -20,13 +20,6 @@ import (
 	"net/netip"
 )
 
-// bitAt returns the bit at the i-th position in the byte slice.
-// The return value is either 0 or 1 as uint8.
-// Panics if the index is out of bounds.
-func bitAt(bytes []byte, i int) uint8 {
-	return bytes[i/8] >> (7 - i%8) & 1
-}
-
 type prefixTreeNode struct {
 	masked bool
 	prefix netip.Prefix

pkg/provider/loadbalancer/iputil/internal/prefixtree_test.go

@@ -28,31 +28,6 @@ import (
 	"sigs.k8s.io/cloud-provider-azure/pkg/provider/loadbalancer/fnutil"
 )
 
-func Test_bitAt(t *testing.T) {
-	bytes := []byte{0b1010_1010, 0b0101_0101}
-	assert.Equal(t, uint8(1), bitAt(bytes, 0))
-	assert.Equal(t, uint8(0), bitAt(bytes, 1))
-	assert.Equal(t, uint8(1), bitAt(bytes, 2))
-	assert.Equal(t, uint8(0), bitAt(bytes, 3))
-
-	assert.Equal(t, uint8(1), bitAt(bytes, 4))
-	assert.Equal(t, uint8(0), bitAt(bytes, 5))
-	assert.Equal(t, uint8(1), bitAt(bytes, 6))
-	assert.Equal(t, uint8(0), bitAt(bytes, 7))
-
-	assert.Equal(t, uint8(0), bitAt(bytes, 8))
-	assert.Equal(t, uint8(1), bitAt(bytes, 9))
-	assert.Equal(t, uint8(0), bitAt(bytes, 10))
-	assert.Equal(t, uint8(1), bitAt(bytes, 11))
-
-	assert.Equal(t, uint8(0), bitAt(bytes, 12))
-	assert.Equal(t, uint8(1), bitAt(bytes, 13))
-	assert.Equal(t, uint8(0), bitAt(bytes, 14))
-	assert.Equal(t, uint8(1), bitAt(bytes, 15))
-
-	assert.Panics(t, func() { bitAt(bytes, 16) })
-}
-
 func TestPrefixTreeIPv4(t *testing.T) {
 	tests := []struct {
 		Name   string

pkg/provider/loadbalancer/iputil/prefix.go

@@ -77,26 +77,29 @@ func GroupPrefixesByFamily(vs []netip.Prefix) ([]netip.Prefix, []netip.Prefix) {
 	return v4, v6
 }
 
-// AggregatePrefixes aggregates prefixes.
-// Overlapping prefixes are merged.
+// AggregatePrefixes merges overlapping or adjacent prefixes.
+// It combines prefixes that can be represented by a single, larger prefix.
+//
+// Examples:
+//   - [192.168.0.0/32, 192.168.0.1/32] -> [192.168.0.0/31] (adjacent)
+//   - [192.168.0.0/24, 192.168.0.1/32] -> [192.168.0.0/24] (overlapping)
+//   - [10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16] -> [10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16] (non-overlapping)
 func AggregatePrefixes(prefixes []netip.Prefix) []netip.Prefix {
 	var (
 		v4, v6 = GroupPrefixesByFamily(prefixes)
-		v4Tree = internal.NewPrefixTreeForIPv4()
-		v6Tree = internal.NewPrefixTreeForIPv6()
 	)
 
-	for _, p := range v4 {
-		v4Tree.Add(p)
-	}
-	for _, p := range v6 {
-		v6Tree.Add(p)
-	}
+	v4 = internal.AggregatePrefixesForSingleIPFamily(v4)
+	v6 = internal.AggregatePrefixesForSingleIPFamily(v6)
 
-	return append(v4Tree.List(), v6Tree.List()...)
+	return append(v4, v6...)
 }
 
 // ExcludePrefixes excludes prefixes from the given prefixes.
+//
+// Examples:
+// - ([192.168.0.0/24], [192.168.0.0/25]) -> [192.168.0.128/25]
+// - ([2001:db8::/64], [2001:db8::1/128, 2001:db8::2/128]) -> [2001:db8::/64]
 func ExcludePrefixes(prefixes []netip.Prefix, exclude []netip.Prefix) []netip.Prefix {
 	var (
 		v4Tree = internal.NewPrefixTreeForIPv4()