/
drainer_util.go
93 lines (80 loc) · 2.63 KB
/
drainer_util.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
package drainer
import (
"github.com/hashicorp/nomad/nomad/structs"
)
const (
// defaultMaxIdsPerTxn is the maximum number of IDs that can be included in a
// single Raft transaction. This is to ensure that the Raft message
// does not become too large.
defaultMaxIdsPerTxn = (1024 * 256) / 36 // 0.25 MB of ids.
)
// partitionIds takes a set of IDs and returns a partitioned view of them such
// that no batch would result in an overly large raft transaction.
func partitionIds(maxIds int, ids []string) [][]string {
index := 0
total := len(ids)
var partitions [][]string
for remaining := total - index; remaining > 0; remaining = total - index {
if remaining < maxIds {
partitions = append(partitions, ids[index:])
break
} else {
partitions = append(partitions, ids[index:index+maxIds])
index += maxIds
}
}
return partitions
}
// transitionTuple is used to group desired transitions and evals
type transitionTuple struct {
Transitions map[string]*structs.DesiredTransition
Evals []*structs.Evaluation
}
// partitionAllocDrain returns a list of alloc transitions and evals to apply
// in a single raft transaction.This is necessary to ensure that the Raft
// transaction does not become too large.
func partitionAllocDrain(maxIds int, transitions map[string]*structs.DesiredTransition,
evals []*structs.Evaluation) []*transitionTuple {
// Determine a stable ordering of the transitioning allocs
allocs := make([]string, 0, len(transitions))
for id := range transitions {
allocs = append(allocs, id)
}
var requests []*transitionTuple
submittedEvals, submittedTrans := 0, 0
for submittedEvals != len(evals) || submittedTrans != len(transitions) {
req := &transitionTuple{
Transitions: make(map[string]*structs.DesiredTransition),
}
requests = append(requests, req)
available := maxIds
// Add the allocs first
if remaining := len(allocs) - submittedTrans; remaining > 0 {
if remaining <= available {
for _, id := range allocs[submittedTrans:] {
req.Transitions[id] = transitions[id]
}
available -= remaining
submittedTrans += remaining
} else {
for _, id := range allocs[submittedTrans : submittedTrans+available] {
req.Transitions[id] = transitions[id]
}
submittedTrans += available
// Exhausted space so skip adding evals
continue
}
}
// Add the evals
if remaining := len(evals) - submittedEvals; remaining > 0 {
if remaining <= available {
req.Evals = evals[submittedEvals:]
submittedEvals += remaining
} else {
req.Evals = evals[submittedEvals : submittedEvals+available]
submittedEvals += available
}
}
}
return requests
}