-
Notifications
You must be signed in to change notification settings - Fork 397
/
state.go
144 lines (115 loc) · 3.77 KB
/
state.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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
// Copyright (C) 2020 Storj Labs, Incache.
// See LICENSE for copying information.
package nodeselection
import (
"context"
"sync"
"github.com/zeebo/errs"
"storj.io/common/storj"
)
// ErrNotEnoughNodes is when selecting nodes failed with the given parameters.
var ErrNotEnoughNodes = errs.Class("not enough nodes")
// State defines a node selector state that allows for selection.
type State struct {
mu sync.RWMutex
stats Stats
// netByID returns subnet based on storj.NodeID
netByID map[storj.NodeID]string
// nonDistinct contains selectors for non-distinct selection.
nonDistinct struct {
Reputable SelectByID
New SelectByID
}
// distinct contains selectors for distinct slection.
distinct struct {
Reputable SelectBySubnet
New SelectBySubnet
}
}
// Stats contains state information.
type Stats struct {
New int
Reputable int
NewDistinct int
ReputableDistinct int
}
// Selector defines interface for selecting nodes.
type Selector interface {
// Count returns the number of maximum number of nodes that it can return.
Count() int
// Select selects up-to n nodes and excluding the IDs.
// When excludedNets is non-nil it will ensure that selected network is unique.
Select(n int, excludedIDs []storj.NodeID, excludeNets map[string]struct{}) []*Node
}
// NewState returns a state based on the input.
func NewState(reputableNodes, newNodes []*Node) *State {
state := &State{}
state.netByID = map[storj.NodeID]string{}
for _, node := range reputableNodes {
state.netByID[node.ID] = node.LastNet
}
for _, node := range newNodes {
state.netByID[node.ID] = node.LastNet
}
state.nonDistinct.Reputable = SelectByID(reputableNodes)
state.nonDistinct.New = SelectByID(newNodes)
state.distinct.Reputable = SelectBySubnetFromNodes(reputableNodes)
state.distinct.New = SelectBySubnetFromNodes(newNodes)
state.stats = Stats{
New: state.nonDistinct.New.Count(),
Reputable: state.nonDistinct.Reputable.Count(),
NewDistinct: state.distinct.New.Count(),
ReputableDistinct: state.distinct.Reputable.Count(),
}
return state
}
// Request contains arguments for State.Request.
type Request struct {
Count int
NewFraction float64
Distinct bool
ExcludedIDs []storj.NodeID
}
// Select selects requestedCount nodes where there will be newFraction nodes.
func (state *State) Select(ctx context.Context, request Request) (_ []*Node, err error) {
defer mon.Task()(&ctx)(&err)
state.mu.RLock()
defer state.mu.RUnlock()
totalCount := request.Count
newCount := int(float64(totalCount) * request.NewFraction)
var selected []*Node
var excludedNets map[string]struct{}
var reputableNodes Selector
var newNodes Selector
if request.Distinct {
excludedNets = map[string]struct{}{}
for _, id := range request.ExcludedIDs {
if net, ok := state.netByID[id]; ok {
excludedNets[net] = struct{}{}
}
}
reputableNodes = state.distinct.Reputable
newNodes = state.distinct.New
} else {
reputableNodes = state.nonDistinct.Reputable
newNodes = state.nonDistinct.New
}
// Get a random selection of new nodes out of the cache first so that if there aren't
// enough new nodes on the network, we can fall back to using reputable nodes instead.
selected = append(selected,
newNodes.Select(newCount, request.ExcludedIDs, excludedNets)...)
// Get all the remaining reputable nodes.
reputableCount := totalCount - len(selected)
selected = append(selected,
reputableNodes.Select(reputableCount, request.ExcludedIDs, excludedNets)...)
if len(selected) < totalCount {
return selected, ErrNotEnoughNodes.New("requested from cache %d, found %d", totalCount, len(selected))
}
return selected, nil
}
// Stats returns state information.
func (state *State) Stats() Stats {
state.mu.RLock()
defer state.mu.RUnlock()
return state.stats
}