d2033c2f52
Currently nodeselection package only contained state for uploads, move these to a subpackage, such that we can make another "downloadselection" for downloads. Then move selection logic from overlay to nodeselection. Change-Id: I0fc42bcae3a29db2728dae9f3863b1e95bf5165b
120 lines
2.7 KiB
Go
120 lines
2.7 KiB
Go
// Copyright (C) 2020 Storj Labs, Incache.
|
|
// See LICENSE for copying information.
|
|
|
|
package uploadselection
|
|
|
|
import (
|
|
mathrand "math/rand" // Using mathrand here because crypto-graphic randomness is not required and simplifies code.
|
|
|
|
"storj.io/common/storj"
|
|
)
|
|
|
|
// SelectByID implements selection from nodes with every node having equal probability.
|
|
type SelectByID []*Node
|
|
|
|
var _ Selector = (SelectByID)(nil)
|
|
|
|
// Count returns the number of maximum number of nodes that it can return.
|
|
func (nodes SelectByID) Count() int { return len(nodes) }
|
|
|
|
// Select selects upto n nodes.
|
|
func (nodes SelectByID) Select(n int, excludedIDs []storj.NodeID, excludedNets map[string]struct{}) []*Node {
|
|
if n <= 0 {
|
|
return nil
|
|
}
|
|
|
|
selected := []*Node{}
|
|
for _, idx := range mathrand.Perm(len(nodes)) {
|
|
node := nodes[idx]
|
|
|
|
if ContainsID(excludedIDs, node.ID) {
|
|
continue
|
|
}
|
|
if excludedNets != nil {
|
|
if _, excluded := excludedNets[node.LastNet]; excluded {
|
|
continue
|
|
}
|
|
excludedNets[node.LastNet] = struct{}{}
|
|
}
|
|
|
|
selected = append(selected, node.Clone())
|
|
if len(selected) >= n {
|
|
break
|
|
}
|
|
}
|
|
|
|
return selected
|
|
}
|
|
|
|
// SelectBySubnet implements selection from nodes with every subnet having equal probability.
|
|
type SelectBySubnet []Subnet
|
|
|
|
var _ Selector = (SelectBySubnet)(nil)
|
|
|
|
// Subnet groups together nodes with the same subnet.
|
|
type Subnet struct {
|
|
Net string
|
|
Nodes []*Node
|
|
}
|
|
|
|
// SelectBySubnetFromNodes creates SelectBySubnet selector from nodes.
|
|
func SelectBySubnetFromNodes(nodes []*Node) SelectBySubnet {
|
|
bynet := map[string][]*Node{}
|
|
for _, node := range nodes {
|
|
bynet[node.LastNet] = append(bynet[node.LastNet], node)
|
|
}
|
|
|
|
var subnets SelectBySubnet
|
|
for net, nodes := range bynet {
|
|
subnets = append(subnets, Subnet{
|
|
Net: net,
|
|
Nodes: nodes,
|
|
})
|
|
}
|
|
|
|
return subnets
|
|
}
|
|
|
|
// Count returns the number of maximum number of nodes that it can return.
|
|
func (subnets SelectBySubnet) Count() int { return len(subnets) }
|
|
|
|
// Select selects upto n nodes.
|
|
func (subnets SelectBySubnet) Select(n int, excludedIDs []storj.NodeID, excludedNets map[string]struct{}) []*Node {
|
|
if n <= 0 {
|
|
return nil
|
|
}
|
|
|
|
selected := []*Node{}
|
|
for _, idx := range mathrand.Perm(len(subnets)) {
|
|
subnet := subnets[idx]
|
|
node := subnet.Nodes[mathrand.Intn(len(subnet.Nodes))]
|
|
|
|
if ContainsID(excludedIDs, node.ID) {
|
|
continue
|
|
}
|
|
if excludedNets != nil {
|
|
if _, excluded := excludedNets[node.LastNet]; excluded {
|
|
continue
|
|
}
|
|
excludedNets[node.LastNet] = struct{}{}
|
|
}
|
|
|
|
selected = append(selected, node.Clone())
|
|
if len(selected) >= n {
|
|
break
|
|
}
|
|
}
|
|
|
|
return selected
|
|
}
|
|
|
|
// ContainsID returns whether ids contains id.
|
|
func ContainsID(ids []storj.NodeID, id storj.NodeID) bool {
|
|
for _, k := range ids {
|
|
if k == id {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|