open-nomad/nomad/drainer/drainer_util.go

97 lines
2.7 KiB
Go

// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0
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
}