forked from NVIDIA/aistore
/
stages.go
71 lines (60 loc) · 1.77 KB
/
stages.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
// Package reb provides global cluster-wide rebalance upon adding/removing storage nodes.
/*
* Copyright (c) 2018-2023, NVIDIA CORPORATION. All rights reserved.
*/
package reb
import (
"sync"
"github.com/artashesbalabekyan/aistore/cluster/meta"
"github.com/artashesbalabekyan/aistore/cmn/atomic"
)
////////////////
// nodeStages //
////////////////
type nodeStages struct {
targets map[string]uint32 // remote tid <-> stage
stage atomic.Uint32 // rebStage* enum: my own current stage
mtx sync.Mutex // updated from different goroutines
}
func newNodeStages() *nodeStages {
return &nodeStages{targets: make(map[string]uint32)}
}
// Returns true if the target is in `newStage` or in any next stage
func (*nodeStages) stageReached(stage, newStage uint32) bool {
return stage > newStage
}
// Mark a 'node' that it has reached the 'stage'. Do nothing if the target
// is already in this stage or has finished it already
func (ns *nodeStages) setStage(daemonID string, stage uint32) {
ns.mtx.Lock()
status, ok := ns.targets[daemonID]
if !ok {
ns.targets[daemonID] = status
}
if !ns.stageReached(status, stage) {
ns.targets[daemonID] = stage
}
ns.mtx.Unlock()
}
// Returns true if the target is in `newStage` or in any next stage.
func (ns *nodeStages) isInStage(si *meta.Snode, stage uint32) bool {
ns.mtx.Lock()
inStage := ns.isInStageUnlocked(si, stage)
ns.mtx.Unlock()
return inStage
}
// Returns true if the target is in `newStage` or in any next stage
func (ns *nodeStages) isInStageUnlocked(si *meta.Snode, stage uint32) bool {
status, ok := ns.targets[si.ID()]
if !ok {
return false
}
return ns.stageReached(status, stage)
}
func (ns *nodeStages) cleanup() {
ns.mtx.Lock()
for k := range ns.targets {
delete(ns.targets, k)
}
ns.mtx.Unlock()
}