Skip to content

Commit 8ac3152

Browse files
zarvdzarvd
authored
Optimize CIDR aggregation to improve performance and reduce memory usage (#7201)
* Optimize CIDR aggregation to improve performance and reduce memory usage * Fix imports
1 parent 02f35e2 commit 8ac3152

File tree

6 files changed

+525
-63
lines changed

6 files changed

+525
-63
lines changed

‎pkg/util/collectionutil/map.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,14 @@ limitations under the License.
1616

1717
package fnutil
1818

19+
func Keys[K comparable, V any](m map[K]V) []K {
20+
rv := make([]K, 0, len(m))
21+
for k := range m {
22+
rv = append(rv, k)
23+
}
24+
return rv
25+
}
26+
1927
func Values[K comparable, V any](m map[K]V) []V {
2028
rv := make([]V, 0, len(m))
2129
for _, v := range m {

‎pkg/util/iputil/bits.go

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/*
2+
Copyright 2024 The Kubernetes Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package iputil
18+
19+
// setBitAt sets the bit at the i-th position in the byte slice to the given value.
20+
// Panics if the index is out of bounds.
21+
// For example,
22+
// - setBitAt([0x00, 0x00], 8, 1) returns [0x00, 0b1000_0000].
23+
// - setBitAt([0xff, 0xff], 0, 0) returns [0b0111_1111, 0xff].
24+
func setBitAt(bytes []byte, i int, bit uint8) {
25+
if bit == 1 {
26+
bytes[i/8] |= 1 << (7 - i%8)
27+
} else {
28+
bytes[i/8] &^= 1 << (7 - i%8)
29+
}
30+
}
31+
32+
// bitAt returns the bit at the i-th position in the byte slice.
33+
// The return value is either 0 or 1 as uint8.
34+
// Panics if the index is out of bounds.
35+
func bitAt(bytes []byte, i int) uint8 {
36+
return bytes[i/8] >> (7 - i%8) & 1
37+
}

‎pkg/util/iputil/bits_test.go

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
/*
2+
Copyright 2024 The Kubernetes Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package iputil
18+
19+
import (
20+
"testing"
21+
22+
"github.com/stretchr/testify/assert"
23+
)
24+
25+
func Test_bitAt(t *testing.T) {
26+
bytes := []byte{0b1010_1010, 0b0101_0101}
27+
assert.Equal(t, uint8(1), bitAt(bytes, 0))
28+
assert.Equal(t, uint8(0), bitAt(bytes, 1))
29+
assert.Equal(t, uint8(1), bitAt(bytes, 2))
30+
assert.Equal(t, uint8(0), bitAt(bytes, 3))
31+
32+
assert.Equal(t, uint8(1), bitAt(bytes, 4))
33+
assert.Equal(t, uint8(0), bitAt(bytes, 5))
34+
assert.Equal(t, uint8(1), bitAt(bytes, 6))
35+
assert.Equal(t, uint8(0), bitAt(bytes, 7))
36+
37+
assert.Equal(t, uint8(0), bitAt(bytes, 8))
38+
assert.Equal(t, uint8(1), bitAt(bytes, 9))
39+
assert.Equal(t, uint8(0), bitAt(bytes, 10))
40+
assert.Equal(t, uint8(1), bitAt(bytes, 11))
41+
42+
assert.Equal(t, uint8(0), bitAt(bytes, 12))
43+
assert.Equal(t, uint8(1), bitAt(bytes, 13))
44+
assert.Equal(t, uint8(0), bitAt(bytes, 14))
45+
assert.Equal(t, uint8(1), bitAt(bytes, 15))
46+
47+
assert.Panics(t, func() { bitAt(bytes, 16) })
48+
}
49+
50+
func Test_setBitAt(t *testing.T) {
51+
tests := []struct {
52+
name string
53+
initial []byte
54+
index int
55+
bit uint8
56+
expected []byte
57+
}{
58+
{
59+
name: "Set first bit to 1",
60+
initial: []byte{0b0000_0000},
61+
index: 0,
62+
bit: 1,
63+
expected: []byte{0b1000_0000},
64+
},
65+
{
66+
name: "Set last bit to 1",
67+
initial: []byte{0b0000_0000},
68+
index: 7,
69+
bit: 1,
70+
expected: []byte{0b0000_0001},
71+
},
72+
{
73+
name: "Set middle bit to 1",
74+
initial: []byte{0b0000_0000},
75+
index: 4,
76+
bit: 1,
77+
expected: []byte{0b0000_1000},
78+
},
79+
{
80+
name: "Set bit to 0",
81+
initial: []byte{0b1111_1111},
82+
index: 3,
83+
bit: 0,
84+
expected: []byte{0b1110_1111},
85+
},
86+
{
87+
name: "Set bit in second byte",
88+
initial: []byte{0b0000_0000, 0b0000_0000},
89+
index: 9,
90+
bit: 1,
91+
expected: []byte{0b0000_0000, 0b0100_0000},
92+
},
93+
}
94+
95+
for _, tt := range tests {
96+
t.Run(tt.name, func(t *testing.T) {
97+
setBitAt(tt.initial, tt.index, tt.bit)
98+
assert.Equal(t, tt.expected, tt.initial)
99+
})
100+
}
101+
102+
assert.Panics(t, func() { setBitAt([]byte{0x00}, 8, 1) })
103+
}

‎pkg/util/iputil/prefix.go

Lines changed: 103 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,10 @@ limitations under the License.
1717
package iputil
1818

1919
import (
20+
"bytes"
2021
"fmt"
2122
"net/netip"
23+
"sort"
2224
)
2325

2426
// IsPrefixesAllowAll returns true if one of the prefixes allows all addresses.
@@ -61,9 +63,108 @@ func GroupPrefixesByFamily(vs []netip.Prefix) ([]netip.Prefix, []netip.Prefix) {
6163
return v4, v6
6264
}
6365

64-
// AggregatePrefixes aggregates prefixes.
65-
// Overlapping prefixes are merged.
66+
// ContainsPrefix checks if prefix p fully contains prefix o.
67+
// It returns true if o is a subset of p, meaning all addresses in o are also in p.
68+
// This is true when p overlaps with o and p has fewer or equal number of bits than o.
69+
func ContainsPrefix(p netip.Prefix, o netip.Prefix) bool {
70+
return p.Bits() <= o.Bits() && p.Overlaps(o)
71+
}
72+
73+
// mergeAdjacentPrefixes attempts to merge two adjacent prefixes into a single prefix.
74+
// It returns the merged prefix and a boolean indicating success.
75+
// Note: This function only merges adjacent prefixes, not overlapping ones.
76+
func mergeAdjacentPrefixes(p1, p2 netip.Prefix) (netip.Prefix, bool) {
77+
// Merge neighboring prefixes if possible
78+
if p1.Bits() != p2.Bits() || p1.Bits() == 0 {
79+
return netip.Prefix{}, false
80+
}
81+
82+
var (
83+
bits = p1.Bits()
84+
p1Bytes = p1.Addr().AsSlice()
85+
p2Bytes = p2.Addr().AsSlice()
86+
)
87+
if bitAt(p1Bytes, bits-1) == 0 {
88+
setBitAt(p1Bytes, bits-1, 1)
89+
} else {
90+
setBitAt(p2Bytes, bits-1, 1)
91+
}
92+
if !bytes.Equal(p1Bytes, p2Bytes) {
93+
return netip.Prefix{}, false
94+
}
95+
96+
rv, _ := p1.Addr().Prefix(bits - 1)
97+
return rv, true
98+
}
99+
100+
// aggregatePrefixesForSingleIPFamily merges overlapping or adjacent prefixes into a single prefix.
101+
// The input prefixes must be the same IP family (IPv4 or IPv6).
102+
// For example,
103+
// - [192.168.0.0/32, 192.168.0.1/32] -> [192.168.0.0/31] (adjacent)
104+
// - [192.168.0.0/24, 192.168.0.1/32] -> [192.168.1.0/24] (overlapping)
105+
func aggregatePrefixesForSingleIPFamily(prefixes []netip.Prefix) []netip.Prefix {
106+
if len(prefixes) <= 1 {
107+
return prefixes
108+
}
109+
110+
sort.Slice(prefixes, func(i, j int) bool {
111+
addrCmp := prefixes[i].Addr().Compare(prefixes[j].Addr())
112+
if addrCmp == 0 {
113+
return prefixes[i].Bits() < prefixes[j].Bits()
114+
}
115+
return addrCmp < 0
116+
})
117+
118+
var rv = []netip.Prefix{prefixes[0]}
119+
120+
for i := 1; i < len(prefixes); i++ {
121+
last, p := rv[len(rv)-1], prefixes[i]
122+
if ContainsPrefix(last, p) {
123+
// Skip overlapping prefixes
124+
continue
125+
}
126+
rv = append(rv, p)
127+
128+
// Merge adjacent prefixes if possible
129+
for len(rv) >= 2 {
130+
// Merge the last two prefixes if they are adjacent
131+
p, ok := mergeAdjacentPrefixes(rv[len(rv)-2], rv[len(rv)-1])
132+
if !ok {
133+
break
134+
}
135+
136+
// Replace the last two prefixes with the merged prefix
137+
rv = rv[:len(rv)-2]
138+
rv = append(rv, p)
139+
}
140+
}
141+
return rv
142+
}
143+
144+
// AggregatePrefixes merges overlapping or adjacent prefixes into a single prefix.
145+
// It combines prefixes that can be represented by a larger, more inclusive prefix.
146+
//
147+
// Examples:
148+
// - Adjacent: [192.168.0.0/32, 192.168.0.1/32] -> [192.168.0.0/31]
149+
// - Overlapping: [192.168.0.0/24, 192.168.0.1/32] -> [192.168.0.0/24]
66150
func AggregatePrefixes(prefixes []netip.Prefix) []netip.Prefix {
151+
var (
152+
v4, v6 = GroupPrefixesByFamily(prefixes)
153+
)
154+
155+
return append(aggregatePrefixesForSingleIPFamily(v4), aggregatePrefixesForSingleIPFamily(v6)...)
156+
}
157+
158+
// AggregatePrefixesWithPrefixTree merges overlapping or adjacent prefixes into a single prefix.
159+
//
160+
// This function uses a prefix tree to aggregate the input prefixes. While it achieves
161+
// the same result as AggregatePrefixes, it is less efficient. For better performance,
162+
// use AggregatePrefixes instead.
163+
//
164+
// Examples:
165+
// - Adjacent: [192.168.0.0/32, 192.168.0.1/32] -> [192.168.0.0/31]
166+
// - Overlapping: [192.168.0.0/24, 192.168.0.1/32] -> [192.168.0.0/24]
167+
func AggregatePrefixesWithPrefixTree(prefixes []netip.Prefix) []netip.Prefix {
67168
var (
68169
v4, v6 = GroupPrefixesByFamily(prefixes)
69170
v4Tree = newPrefixTreeForIPv4()

0 commit comments

Comments
 (0)