Skip to content

Commit 503524c

Browse files
joamakijoamaki
committed
part: Optimize map space usage with one element
In cilium/cilium we're using part.Map in the Backend object to hold the instances of the backend and most of the time it only contains one element. Specialize for this case by adding a singleton field to part.Map. Before (map with 1 element): === RUN TestMapMemoryUse map_test.go:272: 320 bytes per map After: === RUN TestMapMemoryUse map_test.go:348: 40 bytes per map pkg/loadbalancer/benchmark before: Memory statistics from N=10 iterations: Min: Allocated 716852kB in total, 2068061 objects / 125960kB still reachable (per service: 41 objs, 14681B alloc, 2579B in-use) Avg: Allocated 731848kB in total, 2607405 objects / 180868kB still reachable (per service: 52 objs, 14988B alloc, 3704B in-use) Max: Allocated 775441kB in total, 3649005 objects / 317829kB still reachable (per service: 72 objs, 15881B alloc, 6509B in-use) Insert statistics from N=10 iterations: Min: Reconciled 50000 objects in 832.079251ms (16.641µs per service / 60093 services per second) Avg: Reconciled 50000 objects in 883.983229ms (17.679µs per service / 56564 services per second) Max: Reconciled 50000 objects in 988.247623ms (19.764µs per service / 50597 services per second) After: Memory statistics from N=10 iterations: Min: Allocated 669824kB in total, 2067737 objects / 124777kB still reachable (per service: 41 objs, 13718B alloc, 2555B in-use) Avg: Allocated 688314kB in total, 2445133 objects / 166797kB still reachable (per service: 48 objs, 14096B alloc, 3416B in-use) Max: Allocated 734497kB in total, 3444117 objects / 301701kB still reachable (per service: 68 objs, 15042B alloc, 6178B in-use) Insert statistics from N=10 iterations: Min: Reconciled 50000 objects in 798.811656ms (15.976µs per service / 62594 services per second) Avg: Reconciled 50000 objects in 831.351251ms (16.627µs per service / 60143 services per second) Max: Reconciled 50000 objects in 929.973547ms (18.599µs per service / 53766 services per second) Signed-off-by: Jussi Maki <[email protected]>
1 parent d56c21d commit 503524c

File tree

3 files changed

+249
-25
lines changed

3 files changed

+249
-25
lines changed

‎part/map.go‎

Lines changed: 95 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,9 @@ import (
1919
// Map is a typed wrapper around Tree[T] for working with
2020
// keys that are not []byte.
2121
type Map[K, V any] struct {
22-
bytesFromKey func(K) []byte
23-
tree *Tree[mapKVPair[K, V]]
22+
bytesFromKeyFunc func(K) []byte
23+
tree *Tree[mapKVPair[K, V]]
24+
singleton *mapKVPair[K, V]
2425
}
2526

2627
type mapKVPair[K, V any] struct {
@@ -32,14 +33,23 @@ type mapKVPair[K, V any] struct {
3233
// This is not implemented as a method on Map[K,V] as hash maps require the
3334
// comparable constraint and we do not need to limit Map[K, V] to that.
3435
func FromMap[K comparable, V any](m Map[K, V], hm map[K]V) Map[K, V] {
35-
if len(hm) == 0 {
36-
return Map[K, V]{}
36+
switch len(hm) {
37+
case 0:
38+
return m
39+
case 1:
40+
for key, value := range hm {
41+
return m.Set(key, value)
42+
}
3743
}
3844

3945
m.ensureTree()
4046
txn := m.tree.Txn()
41-
for k, v := range hm {
42-
txn.Insert(m.bytesFromKey(k), mapKVPair[K, V]{k, v})
47+
for key, value := range hm {
48+
txn.Insert(m.keyToBytes(key), mapKVPair[K, V]{key, value})
49+
}
50+
if m.singleton != nil {
51+
txn.Insert(m.keyToBytes(m.singleton.Key), *m.singleton)
52+
m.singleton = nil
4353
}
4454
m.tree = txn.CommitOnly()
4555
return m
@@ -52,34 +62,60 @@ func (m *Map[K, V]) ensureTree() {
5262
if m.tree == nil {
5363
m.tree = New[mapKVPair[K, V]](RootOnlyWatch, NoCache)
5464
}
55-
m.bytesFromKey = lookupKeyType[K]()
5665
}
5766

5867
// Get a value from the map by its key.
5968
func (m Map[K, V]) Get(key K) (value V, found bool) {
69+
if m.singleton != nil && bytes.Equal(m.keyToBytes(m.singleton.Key), m.keyToBytes(key)) {
70+
return m.singleton.Value, true
71+
}
72+
6073
if m.tree == nil {
6174
return
6275
}
63-
kv, _, found := m.tree.Get(m.bytesFromKey(key))
76+
kv, _, found := m.tree.Get(m.keyToBytes(key))
6477
return kv.Value, found
6578
}
6679

6780
// Set a value. Returns a new map with the value set.
6881
// Original map is unchanged.
6982
func (m Map[K, V]) Set(key K, value V) Map[K, V] {
83+
keyBytes := m.keyToBytes(key)
84+
if m.tree == nil && m.singleton == nil || m.singleton != nil && bytes.Equal(keyBytes, m.keyToBytes(m.singleton.Key)) {
85+
m.singleton = &mapKVPair[K, V]{key, value}
86+
return m
87+
}
88+
7089
m.ensureTree()
7190
txn := m.tree.Txn()
72-
txn.Insert(m.bytesFromKey(key), mapKVPair[K, V]{key, value})
91+
txn.Insert(m.keyToBytes(key), mapKVPair[K, V]{key, value})
92+
if m.singleton != nil {
93+
txn.Insert(m.keyToBytes(m.singleton.Key), *m.singleton)
94+
m.singleton = nil
95+
}
7396
m.tree = txn.CommitOnly()
7497
return m
7598
}
7699

100+
func (m *Map[K, V]) keyToBytes(key K) []byte {
101+
if m.bytesFromKeyFunc == nil {
102+
m.bytesFromKeyFunc = lookupKeyType[K]()
103+
}
104+
return m.bytesFromKeyFunc(key)
105+
}
106+
77107
// Delete a value from the map. Returns a new map
78108
// without the element pointed to by the key (if found).
79109
func (m Map[K, V]) Delete(key K) Map[K, V] {
110+
if m.singleton != nil {
111+
if bytes.Equal(m.keyToBytes(m.singleton.Key), m.keyToBytes(key)) {
112+
m.singleton = nil
113+
}
114+
return m
115+
}
80116
if m.tree != nil {
81117
txn := m.tree.Txn()
82-
txn.Delete(m.bytesFromKey(key))
118+
txn.Delete(m.keyToBytes(key))
83119
// Map is a struct passed by value, so we can modify
84120
// it without changing the caller's view of it.
85121
m.tree = txn.CommitOnly()
@@ -107,26 +143,47 @@ func toSeq2[K, V any](iter *Iterator[mapKVPair[K, V]]) iter.Seq2[K, V] {
107143
// LowerBound iterates over all keys in order with value equal
108144
// to or greater than [from].
109145
func (m Map[K, V]) LowerBound(from K) iter.Seq2[K, V] {
146+
if m.singleton != nil {
147+
if bytes.Compare(m.keyToBytes(m.singleton.Key), m.keyToBytes(from)) >= 0 {
148+
return m.singletonIter()
149+
}
150+
}
110151
if m.tree == nil {
111152
return toSeq2[K, V](nil)
112153
}
113-
return toSeq2(m.tree.LowerBound(m.bytesFromKey(from)))
154+
return toSeq2(m.tree.LowerBound(m.keyToBytes(from)))
155+
}
156+
157+
func (m *Map[K, V]) singletonIter() iter.Seq2[K, V] {
158+
return func(yield func(K, V) bool) {
159+
if m.singleton != nil {
160+
yield(m.singleton.Key, m.singleton.Value)
161+
}
162+
}
114163
}
115164

116165
// Prefix iterates in order over all keys that start with
117166
// the given prefix.
118167
func (m Map[K, V]) Prefix(prefix K) iter.Seq2[K, V] {
168+
if m.singleton != nil {
169+
if bytes.HasPrefix(m.keyToBytes(m.singleton.Key), m.keyToBytes(prefix)) {
170+
return m.singletonIter()
171+
}
172+
}
119173
if m.tree == nil {
120174
return toSeq2[K, V](nil)
121175
}
122-
iter, _ := m.tree.Prefix(m.bytesFromKey(prefix))
176+
iter, _ := m.tree.Prefix(m.keyToBytes(prefix))
123177
return toSeq2(iter)
124178
}
125179

126180
// All iterates every key-value in the map in order.
127181
// The order is in bytewise order of the byte slice
128182
// returned by bytesFromKey.
129183
func (m Map[K, V]) All() iter.Seq2[K, V] {
184+
if m.singleton != nil {
185+
return m.singletonIter()
186+
}
130187
if m.tree == nil {
131188
return toSeq2[K, V](nil)
132189
}
@@ -136,10 +193,12 @@ func (m Map[K, V]) All() iter.Seq2[K, V] {
136193
// EqualKeys returns true if both maps contain the same keys.
137194
func (m Map[K, V]) EqualKeys(other Map[K, V]) bool {
138195
switch {
139-
case m.tree == nil && other.tree == nil:
140-
return true
141196
case m.Len() != other.Len():
142197
return false
198+
case m.singleton != nil && other.singleton != nil:
199+
return bytes.Equal(m.keyToBytes(m.singleton.Key), other.keyToBytes(other.singleton.Key))
200+
case m.tree == nil && other.tree == nil:
201+
return true
143202
default:
144203
iter1 := m.tree.Iterator()
145204
iter2 := other.tree.Iterator()
@@ -163,10 +222,13 @@ func (m Map[K, V]) EqualKeys(other Map[K, V]) bool {
163222
// slow and mostly useful for testing.
164223
func (m Map[K, V]) SlowEqual(other Map[K, V]) bool {
165224
switch {
166-
case m.tree == nil && other.tree == nil:
167-
return true
168225
case m.Len() != other.Len():
169226
return false
227+
case m.singleton != nil && other.singleton != nil:
228+
return bytes.Equal(m.keyToBytes(m.singleton.Key), other.keyToBytes(other.singleton.Key)) &&
229+
reflect.DeepEqual(m.singleton.Value, other.singleton.Value)
230+
case m.tree == nil && other.tree == nil:
231+
return true
170232
default:
171233
iter1 := m.tree.Iterator()
172234
iter2 := other.tree.Iterator()
@@ -187,6 +249,9 @@ func (m Map[K, V]) SlowEqual(other Map[K, V]) bool {
187249

188250
// Len returns the number of elements in the map.
189251
func (m Map[K, V]) Len() int {
252+
if m.singleton != nil {
253+
return 1
254+
}
190255
if m.tree == nil {
191256
return 0
192257
}
@@ -218,6 +283,8 @@ func (m Map[K, V]) MarshalJSON() ([]byte, error) {
218283
}
219284

220285
func (m *Map[K, V]) UnmarshalJSON(data []byte) error {
286+
*m = Map[K, V]{}
287+
221288
dec := json.NewDecoder(bytes.NewReader(data))
222289
t, err := dec.Token()
223290
if err != nil {
@@ -234,7 +301,7 @@ func (m *Map[K, V]) UnmarshalJSON(data []byte) error {
234301
if err != nil {
235302
return err
236303
}
237-
txn.Insert(m.bytesFromKey(kv.Key), mapKVPair[K, V]{kv.Key, kv.Value})
304+
txn.Insert(m.keyToBytes(kv.Key), mapKVPair[K, V]{kv.Key, kv.Value})
238305
}
239306

240307
t, err = dec.Token()
@@ -245,21 +312,26 @@ func (m *Map[K, V]) UnmarshalJSON(data []byte) error {
245312
return fmt.Errorf("%T.UnmarshalJSON: expected ']' got %v", m, t)
246313
}
247314
m.tree = txn.CommitOnly()
315+
316+
if m.tree.size == 1 {
317+
_, kv, _ := m.tree.Iterator().Next()
318+
m.singleton = &kv
319+
m.tree = nil
320+
}
248321
return nil
249322
}
250323

251324
func (m Map[K, V]) MarshalYAML() (any, error) {
252325
kvs := make([]mapKVPair[K, V], 0, m.Len())
253-
if m.tree != nil {
254-
iter := m.tree.Iterator()
255-
for _, kv, ok := iter.Next(); ok; _, kv, ok = iter.Next() {
256-
kvs = append(kvs, kv)
257-
}
326+
for k, v := range m.All() {
327+
kvs = append(kvs, mapKVPair[K, V]{k, v})
258328
}
259329
return kvs, nil
260330
}
261331

262332
func (m *Map[K, V]) UnmarshalYAML(value *yaml.Node) error {
333+
*m = Map[K, V]{}
334+
263335
if value.Kind != yaml.SequenceNode {
264336
return fmt.Errorf("%T.UnmarshalYAML: expected sequence", m)
265337
}
@@ -273,7 +345,7 @@ func (m *Map[K, V]) UnmarshalYAML(value *yaml.Node) error {
273345
if err := e.Decode(&kv); err != nil {
274346
return err
275347
}
276-
txn.Insert(m.bytesFromKey(kv.Key), mapKVPair[K, V]{kv.Key, kv.Value})
348+
txn.Insert(m.keyToBytes(kv.Key), mapKVPair[K, V]{kv.Key, kv.Value})
277349
}
278350
m.tree = txn.CommitOnly()
279351
return nil

‎part/map_test.go‎

Lines changed: 100 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"encoding/json"
88
"fmt"
99
"iter"
10+
"maps"
1011
"math/rand/v2"
1112
"runtime"
1213
"testing"
@@ -55,12 +56,14 @@ func TestStringMap(t *testing.T) {
5556

5657
// Set some values in two different ways.
5758
m = m.Set("1_one", 1)
59+
assert.Equal(t, 1, m.Len())
5860
m = part.FromMap(m, map[string]int{
5961
"2_two": 2,
6062
"3_three": 3,
6163
})
64+
assert.Equal(t, 3, m.Len())
6265

63-
// Setting on a copy doeen't affect original
66+
// Setting on a copy doesn't affect original
6467
m.Set("4_four", 4)
6568
_, ok = m.Get("4_four")
6669
assert.False(t, ok, "Get non-existing")
@@ -78,6 +81,7 @@ func TestStringMap(t *testing.T) {
7881

7982
expected := kvs
8083
for k, v := range m.All() {
84+
t.Logf("%v %v", k, v)
8185
kv := expected[0]
8286
expected = expected[1:]
8387
assert.EqualValues(t, kv.k, k)
@@ -117,6 +121,100 @@ func TestStringMap(t *testing.T) {
117121
assert.False(t, ok, "Get after Delete")
118122

119123
assert.Equal(t, 2, m.Len())
124+
125+
var m3 part.Map[string, int]
126+
bs, err := m.MarshalJSON()
127+
assert.NoError(t, err)
128+
assert.NoError(t, m3.UnmarshalJSON(bs))
129+
assert.Equal(t, 2, m3.Len())
130+
assert.True(t, m.SlowEqual(m3))
131+
132+
m3 = part.Map[string, int]{}
133+
bs, err = yaml.Marshal(m)
134+
assert.NoError(t, err)
135+
assert.NoError(t, yaml.Unmarshal(bs, &m3))
136+
assert.Equal(t, 2, m3.Len())
137+
assert.True(t, m.SlowEqual(m3))
138+
}
139+
140+
func TestSingletonMap(t *testing.T) {
141+
var m part.Map[string, int]
142+
m = m.Set("one", 1)
143+
assert.Equal(t, 1, m.Len())
144+
assert.False(t, m.SlowEqual(part.Map[string, int]{}))
145+
assert.True(t, m.SlowEqual(m))
146+
147+
v, found := m.Get("nope")
148+
assert.False(t, found)
149+
assert.Zero(t, v)
150+
151+
v, found = m.Get("one")
152+
assert.True(t, found)
153+
assert.Equal(t, 1, v)
154+
155+
m2 := m.Set("one", 2)
156+
v, found = m.Get("one")
157+
assert.True(t, found)
158+
assert.Equal(t, 1, v)
159+
v, found = m2.Get("one")
160+
assert.True(t, found)
161+
assert.Equal(t, 2, v)
162+
assert.True(t, m.EqualKeys(m2))
163+
assert.True(t, m2.EqualKeys(m))
164+
assert.False(t, m.SlowEqual(m2))
165+
assert.False(t, m2.SlowEqual(m))
166+
assert.True(t, m2.SlowEqual(m2))
167+
m2 = m2.Delete("nope")
168+
m2 = m2.Delete("one")
169+
assert.Equal(t, 0, m2.Len())
170+
_, found = m2.Get("one")
171+
assert.False(t, found)
172+
assert.False(t, m.EqualKeys(m2))
173+
assert.False(t, m2.EqualKeys(m))
174+
assert.False(t, m.SlowEqual(m2))
175+
assert.False(t, m2.SlowEqual(m))
176+
assert.Equal(t, 0, m2.Len())
177+
178+
x := maps.Collect(m.Prefix(""))
179+
assert.Equal(t, 1, x["one"])
180+
x = maps.Collect(m.Prefix("o"))
181+
assert.Equal(t, 1, x["one"])
182+
x = maps.Collect(m.Prefix("one"))
183+
assert.Equal(t, 1, x["one"])
184+
x = maps.Collect(m.Prefix("one1"))
185+
assert.Len(t, x, 0)
186+
187+
x = maps.Collect(m.LowerBound(""))
188+
assert.Equal(t, 1, x["one"])
189+
x = maps.Collect(m.LowerBound("a"))
190+
assert.Equal(t, 1, x["one"])
191+
x = maps.Collect(m.LowerBound("one"))
192+
assert.Equal(t, 1, x["one"])
193+
x = maps.Collect(m.LowerBound("one1"))
194+
assert.Len(t, x, 0)
195+
196+
m2 = part.Map[string, int]{}
197+
m2 = part.FromMap(m2, nil)
198+
assert.Equal(t, 0, m2.Len())
199+
200+
m2 = part.FromMap(m, nil)
201+
assert.True(t, m.SlowEqual(m2))
202+
assert.True(t, m2.SlowEqual(m))
203+
204+
m2 = part.FromMap(m, map[string]int{"one": 2})
205+
assert.Equal(t, 1, m2.Len())
206+
v, found = m2.Get("one")
207+
assert.True(t, found)
208+
assert.Equal(t, 2, v)
209+
210+
m2 = part.FromMap(m2, map[string]int{"two": 2})
211+
assert.Equal(t, 2, m2.Len())
212+
v, found = m2.Get("one")
213+
assert.True(t, found)
214+
assert.Equal(t, 2, v)
215+
v, found = m2.Get("two")
216+
assert.True(t, found)
217+
assert.Equal(t, 2, v)
120218
}
121219

122220
func TestUint64Map(t *testing.T) {
@@ -274,7 +372,7 @@ func TestMapMemoryUse(t *testing.T) {
274372
// Do some thing with the maps to ensure they weren't GCd.
275373
for _, m := range maps {
276374
if m.Len() != 1 {
277-
t.Fatalf("bad count")
375+
t.Fatalf("bad count %d", m.Len())
278376
}
279377
}
280378
}

0 commit comments

Comments
 (0)