github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/requester/ranking/chain.go (about)

     1  package ranking
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/filecoin-project/bacalhau/pkg/model"
     7  	"github.com/filecoin-project/bacalhau/pkg/requester"
     8  	"github.com/libp2p/go-libp2p/core/peer"
     9  )
    10  
    11  // Chain assigns a random rank to each node to allow the requester to select random top nodes
    12  // for job execution.
    13  type Chain struct {
    14  	rankers []requester.NodeRanker
    15  }
    16  
    17  func NewChain() *Chain {
    18  	return &Chain{}
    19  }
    20  
    21  // Add ranker to the chain
    22  func (c *Chain) Add(ranker ...requester.NodeRanker) {
    23  	c.rankers = append(c.rankers, ranker...)
    24  }
    25  
    26  func (c *Chain) RankNodes(ctx context.Context, job model.Job, nodes []model.NodeInfo) ([]requester.NodeRank, error) {
    27  	// initialize map of node ranks
    28  	ranksMap := make(map[peer.ID]*requester.NodeRank, len(nodes))
    29  	for _, node := range nodes {
    30  		ranksMap[node.PeerInfo.ID] = &requester.NodeRank{NodeInfo: node, Rank: 0}
    31  	}
    32  
    33  	// iterate over the rankers and add their ranks to the map
    34  	// once a node is ranked below zero, it is not considered for job execution and the rank will never be increased above zero
    35  	// by other rankers. It can only go down more
    36  	for _, ranker := range c.rankers {
    37  		nodeRanks, err := ranker.RankNodes(ctx, job, nodes)
    38  		if err != nil {
    39  			return nil, err
    40  		}
    41  		for _, nodeRank := range nodeRanks {
    42  			if ranksMap[nodeRank.NodeInfo.PeerInfo.ID].Rank < 0 || nodeRank.Rank < 0 {
    43  				ranksMap[nodeRank.NodeInfo.PeerInfo.ID].Rank = -1
    44  			} else {
    45  				ranksMap[nodeRank.NodeInfo.PeerInfo.ID].Rank += nodeRank.Rank
    46  			}
    47  		}
    48  	}
    49  
    50  	nodeRanks := make([]requester.NodeRank, 0, len(ranksMap))
    51  	for _, nodeRank := range ranksMap {
    52  		nodeRanks = append(nodeRanks, *nodeRank)
    53  	}
    54  	return nodeRanks, nil
    55  }