/
invariant.go
60 lines (50 loc) · 1.62 KB
/
invariant.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
// Copyright (C) 2023 Storj Labs, Inc.
// See LICENSE for copying information.
package nodeselection
import (
"storj.io/storj/private/intset"
"storj.io/storj/satellite/metabase"
)
// Invariant checks the current placement, and identifies the pieces which should be moved.
// Used by repair jobs.
type Invariant func(pieces metabase.Pieces, nodes []SelectedNode) intset.Set
// AllGood is an invariant, which accepts all piece sets as good.
func AllGood() Invariant {
return func(pieces metabase.Pieces, nodes []SelectedNode) intset.Set {
return intset.NewSet(0)
}
}
// ClumpingByAttribute allows only one selected piece by attribute groups.
func ClumpingByAttribute(attr NodeAttribute, maxAllowed int) Invariant {
return func(pieces metabase.Pieces, nodes []SelectedNode) intset.Set {
usedGroups := make(map[string]int, len(pieces))
maxPieceNum := 0
for _, piece := range pieces {
if int(piece.Number) > maxPieceNum {
maxPieceNum = int(piece.Number)
}
}
maxPieceNum++
res := intset.NewSet(maxPieceNum)
for index, nodeRecord := range nodes {
attribute := attr(nodeRecord)
if attribute == "" {
continue
}
pieceNum := pieces[index].Number
count := usedGroups[attribute]
if count >= maxAllowed {
// this group was already seen, enough times
res.Include(int(pieceNum))
} else {
// add to the list of seen groups
usedGroups[attribute] = count + 1
}
}
return res
}
}
// ClumpingByAnyTag tries to limit the number of nodes with the same tag value.
func ClumpingByAnyTag(key string, maxAllowed int) Invariant {
return ClumpingByAttribute(AnyNodeTagAttribute(key), maxAllowed)
}