github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/reb/stages.go (about)

     1  // Package reb provides global cluster-wide rebalance upon adding/removing storage nodes.
     2  /*
     3   * Copyright (c) 2018-2023, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package reb
     6  
     7  import (
     8  	"sync"
     9  
    10  	"github.com/NVIDIA/aistore/cmn/atomic"
    11  	"github.com/NVIDIA/aistore/core/meta"
    12  )
    13  
    14  ////////////////
    15  // nodeStages //
    16  ////////////////
    17  
    18  type nodeStages struct {
    19  	targets map[string]uint32 // remote tid <-> stage
    20  	stage   atomic.Uint32     // rebStage* enum: my own current stage
    21  	mtx     sync.Mutex        // updated from different goroutines
    22  }
    23  
    24  func newNodeStages() *nodeStages {
    25  	return &nodeStages{targets: make(map[string]uint32)}
    26  }
    27  
    28  // Returns true if the target is in `newStage` or in any next stage
    29  func (*nodeStages) stageReached(stage, newStage uint32) bool {
    30  	return stage > newStage
    31  }
    32  
    33  // Mark a 'node' that it has reached the 'stage'. Do nothing if the target
    34  // is already in this stage or has finished it already
    35  func (ns *nodeStages) setStage(daemonID string, stage uint32) {
    36  	ns.mtx.Lock()
    37  	status, ok := ns.targets[daemonID]
    38  	if !ok {
    39  		ns.targets[daemonID] = status
    40  	}
    41  
    42  	if !ns.stageReached(status, stage) {
    43  		ns.targets[daemonID] = stage
    44  	}
    45  	ns.mtx.Unlock()
    46  }
    47  
    48  // Returns true if the target is in `newStage` or in any next stage.
    49  func (ns *nodeStages) isInStage(si *meta.Snode, stage uint32) bool {
    50  	ns.mtx.Lock()
    51  	inStage := ns.isInStageUnlocked(si, stage)
    52  	ns.mtx.Unlock()
    53  	return inStage
    54  }
    55  
    56  // Returns true if the target is in `newStage` or in any next stage
    57  func (ns *nodeStages) isInStageUnlocked(si *meta.Snode, stage uint32) bool {
    58  	status, ok := ns.targets[si.ID()]
    59  	if !ok {
    60  		return false
    61  	}
    62  	return ns.stageReached(status, stage)
    63  }
    64  
    65  func (ns *nodeStages) cleanup() {
    66  	ns.mtx.Lock()
    67  	clear(ns.targets)
    68  	ns.mtx.Unlock()
    69  }