/
sample.go
105 lines (91 loc) · 2.76 KB
/
sample.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
// SPDX-FileCopyrightText: 2023 SAP SE or an SAP affiliate company and Gardener contributors
//
// SPDX-License-Identifier: Apache-2.0
package config
import (
"math/rand"
"sort"
"sync"
)
// SampleConfig holds configuration and state for node sample.
type SampleConfig struct {
// MaxNodes is the maximum number of sample nodes. If 0, all nodes are selected
MaxNodes int
// NodeSampleStore stores node hostnames with floating index for stable sample selection
NodeSampleStore *NodeSampleStore
}
// NewNodeSampleStore create a new node sample store.
func NewNodeSampleStore(nodeName string) *NodeSampleStore {
return &NodeSampleStore{nodeName: nodeName, store: map[string]float64{}}
}
type orderedNode struct {
hostname string
index float64
}
// NodeSampleStore maps all known hostnames to a random number [0.0, 1.0).
// It allows to keep node samples as stable as possible after adding or removing nodes.
type NodeSampleStore struct {
sync.Mutex
nodeName string
store map[string]float64
}
// SelectTopNodes selects a stable nodes sample of the given size.
func (s *NodeSampleStore) SelectTopNodes(hostnames map[string]struct{}, size int) map[string]struct{} {
s.Lock()
defer s.Unlock()
for name := range s.store {
if _, ok := hostnames[name]; !ok {
delete(s.store, name)
}
}
array := make([]orderedNode, 0, len(hostnames))
for name := range hostnames {
index, ok := s.store[name]
if !ok {
index = rand.Float64()
if name == s.nodeName {
index = 0 // always include own node for self-checks
}
s.store[name] = index
}
array = append(array, orderedNode{hostname: name, index: index})
}
sort.Slice(array, func(i, j int) bool {
return array[i].index < array[j].index
})
if len(array) > size {
array = array[:size]
}
topNodes := map[string]struct{}{}
for _, item := range array {
topNodes[item.hostname] = struct{}{}
}
return topNodes
}
// ShuffledSample selects a node sample and shuffles its order.
func (sc *SampleConfig) ShuffledSample(cc ClusterConfig) ClusterConfig {
return ClusterConfig{
NodeCount: len(cc.Nodes),
Nodes: CloneAndShuffle(selectSample(sc, cc.Nodes)),
PodEndpoints: CloneAndShuffle(selectSample(sc, cc.PodEndpoints)),
InternalKubeAPIServer: cc.InternalKubeAPIServer,
KubeAPIServer: cc.KubeAPIServer,
}
}
func selectSample[T WithDestHost](pc *SampleConfig, items []T) []T {
if pc.MaxNodes == 0 || pc.MaxNodes > len(items) {
return items
}
hostnames := map[string]struct{}{}
for _, item := range items {
hostnames[item.DestHost()] = struct{}{}
}
topNodes := pc.NodeSampleStore.SelectTopNodes(hostnames, pc.MaxNodes)
var sample []T
for _, item := range items {
if _, ok := topNodes[item.DestHost()]; ok {
sample = append(sample, item)
}
}
return sample
}