-
Notifications
You must be signed in to change notification settings - Fork 2k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Backport of drainer: test refactoring to clarify behavior around dele…
…te/down nodes into release/1.5.x (#16622) This pull request was automerged via backport-assistant
- Loading branch information
1 parent
b11eeb4
commit e865333
Showing
4 changed files
with
346 additions
and
171 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,52 +1,144 @@ | ||
| package drainer | ||
|
|
||
| import ( | ||
| "context" | ||
| "sync" | ||
| "testing" | ||
| "time" | ||
|
|
||
| "golang.org/x/time/rate" | ||
|
|
||
| "github.com/hashicorp/nomad/helper/testlog" | ||
| "github.com/hashicorp/nomad/nomad/state" | ||
| "github.com/hashicorp/nomad/nomad/structs" | ||
| ) | ||
|
|
||
| type MockNodeTrackerEvent struct { | ||
| NodeUpdate *structs.Node | ||
| NodeRemove string | ||
| } | ||
| // This file contains helpers for testing. The raft shims make it hard to test | ||
| // the whole package behavior of the drainer. See also nomad/drainer_int_test.go | ||
| // for integration tests. | ||
|
|
||
| type MockNodeTracker struct { | ||
| Nodes map[string]*structs.Node | ||
| Events []*MockNodeTrackerEvent | ||
| type MockJobWatcher struct { | ||
| drainCh chan *DrainRequest | ||
| migratedCh chan []*structs.Allocation | ||
| jobs map[structs.NamespacedID]struct{} | ||
| sync.Mutex | ||
| } | ||
|
|
||
| func NewMockNodeTracker() *MockNodeTracker { | ||
| return &MockNodeTracker{ | ||
| Nodes: make(map[string]*structs.Node), | ||
| Events: make([]*MockNodeTrackerEvent, 0, 16), | ||
| // RegisterJobs marks the job as being watched | ||
| func (m *MockJobWatcher) RegisterJobs(jobs []structs.NamespacedID) { | ||
| m.Lock() | ||
| defer m.Unlock() | ||
| for _, job := range jobs { | ||
| m.jobs[job] = struct{}{} | ||
| } | ||
| } | ||
|
|
||
| func (m *MockNodeTracker) TrackedNodes() map[string]*structs.Node { | ||
| m.Lock() | ||
| defer m.Unlock() | ||
| return m.Nodes | ||
| // Drain returns the DrainRequest channel. Tests can send on this channel to | ||
| // simulate steps through the NodeDrainer watch loop. (Sending on this channel | ||
| // will block anywhere else.) | ||
| func (m *MockJobWatcher) Drain() <-chan *DrainRequest { | ||
| return m.drainCh | ||
| } | ||
|
|
||
| func (m *MockNodeTracker) Remove(nodeID string) { | ||
| m.Lock() | ||
| defer m.Unlock() | ||
| delete(m.Nodes, nodeID) | ||
| m.Events = append(m.Events, &MockNodeTrackerEvent{NodeRemove: nodeID}) | ||
| // Migrated returns the channel of migrated allocations. Tests can send on this | ||
| // channel to simulate steps through the NodeDrainer watch loop. (Sending on | ||
| // this channel will block anywhere else.) | ||
| func (m *MockJobWatcher) Migrated() <-chan []*structs.Allocation { | ||
| return m.migratedCh | ||
| } | ||
|
|
||
| func (m *MockNodeTracker) Update(node *structs.Node) { | ||
| type MockDeadlineNotifier struct { | ||
| expiredCh <-chan []string | ||
| nodes map[string]struct{} | ||
| sync.Mutex | ||
| } | ||
|
|
||
| // NextBatch returns the channel of expired nodes. Tests can send on this | ||
| // channel to simulate timer events in the NodeDrainer watch loop. (Sending on | ||
| // this channel will block anywhere else.) | ||
| func (m *MockDeadlineNotifier) NextBatch() <-chan []string { | ||
| return m.expiredCh | ||
| } | ||
|
|
||
| // Remove removes the given node from being tracked for a deadline. | ||
| func (m *MockDeadlineNotifier) Remove(nodeID string) { | ||
| m.Lock() | ||
| defer m.Unlock() | ||
| m.Nodes[node.ID] = node | ||
| m.Events = append(m.Events, &MockNodeTrackerEvent{NodeUpdate: node}) | ||
| delete(m.nodes, nodeID) | ||
| } | ||
|
|
||
| func (m *MockNodeTracker) events() []*MockNodeTrackerEvent { | ||
| // Watch marks the node as being watched; this mock throws out the timer in lieu | ||
| // of manully sending on the channel to avoid racy tests. | ||
| func (m *MockDeadlineNotifier) Watch(nodeID string, _ time.Time) { | ||
| m.Lock() | ||
| defer m.Unlock() | ||
| m.nodes[nodeID] = struct{}{} | ||
| } | ||
|
|
||
| type MockRaftApplierShim struct { | ||
| lock sync.Mutex | ||
| state *state.StateStore | ||
| } | ||
|
|
||
| // AllocUpdateDesiredTransition mocks a write to raft as a state store update | ||
| func (m *MockRaftApplierShim) AllocUpdateDesiredTransition( | ||
| allocs map[string]*structs.DesiredTransition, evals []*structs.Evaluation) (uint64, error) { | ||
|
|
||
| m.lock.Lock() | ||
| defer m.lock.Unlock() | ||
|
|
||
| index, _ := m.state.LatestIndex() | ||
| index++ | ||
| err := m.state.UpdateAllocsDesiredTransitions(structs.MsgTypeTestSetup, index, allocs, evals) | ||
| return index, err | ||
| } | ||
|
|
||
| // NodesDrainComplete mocks a write to raft as a state store update | ||
| func (m *MockRaftApplierShim) NodesDrainComplete( | ||
| nodes []string, event *structs.NodeEvent) (uint64, error) { | ||
|
|
||
| m.lock.Lock() | ||
| defer m.lock.Unlock() | ||
|
|
||
| index, _ := m.state.LatestIndex() | ||
| index++ | ||
|
|
||
| updates := make(map[string]*structs.DrainUpdate, len(nodes)) | ||
| nodeEvents := make(map[string]*structs.NodeEvent, len(nodes)) | ||
| update := &structs.DrainUpdate{} | ||
| for _, node := range nodes { | ||
| updates[node] = update | ||
| if event != nil { | ||
| nodeEvents[node] = event | ||
| } | ||
| } | ||
| now := time.Now().Unix() | ||
|
|
||
| err := m.state.BatchUpdateNodeDrain(structs.MsgTypeTestSetup, index, now, | ||
| updates, nodeEvents) | ||
|
|
||
| return index, err | ||
| } | ||
|
|
||
| func testNodeDrainWatcher(t *testing.T) (*nodeDrainWatcher, *state.StateStore, *NodeDrainer) { | ||
| t.Helper() | ||
| store := state.TestStateStore(t) | ||
| limiter := rate.NewLimiter(100.0, 100) | ||
| logger := testlog.HCLogger(t) | ||
|
|
||
| drainer := &NodeDrainer{ | ||
| enabled: false, | ||
| logger: logger, | ||
| nodes: map[string]*drainingNode{}, | ||
| jobWatcher: &MockJobWatcher{jobs: map[structs.NamespacedID]struct{}{}}, | ||
| deadlineNotifier: &MockDeadlineNotifier{nodes: map[string]struct{}{}}, | ||
| state: store, | ||
| queryLimiter: limiter, | ||
| raft: &MockRaftApplierShim{state: store}, | ||
| batcher: allocMigrateBatcher{}, | ||
| } | ||
|
|
||
| return m.Events | ||
| w := NewNodeDrainWatcher(context.Background(), limiter, store, logger, drainer) | ||
| drainer.nodeWatcher = w | ||
| return w, store, drainer | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.