diff --git a/satellite/nodeselection/uploadselection/common.go b/satellite/nodeselection/uploadselection/common.go index b682dea6c..4fa1e1b8b 100644 --- a/satellite/nodeselection/uploadselection/common.go +++ b/satellite/nodeselection/uploadselection/common.go @@ -1,4 +1,4 @@ -// Copyright (C) 2020 Storj Labs, Incache. +// Copyright (C) 2020 Storj Labs, Inc. // See LICENSE for copying information. // Package uploadselection implements node selection logic for uploads. diff --git a/satellite/nodeselection/uploadselection/criteria.go b/satellite/nodeselection/uploadselection/criteria.go new file mode 100644 index 000000000..45ab19681 --- /dev/null +++ b/satellite/nodeselection/uploadselection/criteria.go @@ -0,0 +1,36 @@ +// Copyright (C) 2021 Storj Labs, Inc. +// See LICENSE for copying information + +package uploadselection + +import "storj.io/common/storj" + +// Criteria to filter nodes. +type Criteria struct { + ExcludeNodeIDs []storj.NodeID + AutoExcludeSubnets map[string]struct{} // initialize it with empty map to keep only one node per subnet. +} + +// MatchInclude returns with true if node is selected. +func (c *Criteria) MatchInclude(node *Node) bool { + if ContainsID(c.ExcludeNodeIDs, node.ID) { + return false + } + if c.AutoExcludeSubnets != nil { + if _, excluded := c.AutoExcludeSubnets[node.LastNet]; excluded { + return false + } + c.AutoExcludeSubnets[node.LastNet] = struct{}{} + } + return true +} + +// ContainsID returns whether ids contain id. +func ContainsID(ids []storj.NodeID, id storj.NodeID) bool { + for _, k := range ids { + if k == id { + return true + } + } + return false +} diff --git a/satellite/nodeselection/uploadselection/criteria_test.go b/satellite/nodeselection/uploadselection/criteria_test.go new file mode 100644 index 000000000..f382bb755 --- /dev/null +++ b/satellite/nodeselection/uploadselection/criteria_test.go @@ -0,0 +1,83 @@ +// Copyright (C) 2021 Storj Labs, Inc. +// See LICENSE for copying information + +package uploadselection + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "storj.io/common/storj" + "storj.io/common/testrand" +) + +func TestCriteria_AutoExcludeSubnet(t *testing.T) { + + criteria := Criteria{ + AutoExcludeSubnets: map[string]struct{}{}, + } + + assert.True(t, criteria.MatchInclude(&Node{ + LastNet: "192.168.0.1", + })) + + assert.False(t, criteria.MatchInclude(&Node{ + LastNet: "192.168.0.1", + })) + + assert.True(t, criteria.MatchInclude(&Node{ + LastNet: "192.168.1.1", + })) +} + +func TestCriteria_ExcludeNodeID(t *testing.T) { + included := testrand.NodeID() + excluded := testrand.NodeID() + + criteria := Criteria{ + ExcludeNodeIDs: []storj.NodeID{excluded}, + } + + assert.False(t, criteria.MatchInclude(&Node{ + NodeURL: storj.NodeURL{ + ID: excluded, + Address: "localhost", + }, + })) + + assert.True(t, criteria.MatchInclude(&Node{ + NodeURL: storj.NodeURL{ + ID: included, + Address: "localhost", + }, + })) + +} + +func TestCriteria_NodeIDAndSubnet(t *testing.T) { + excluded := testrand.NodeID() + + criteria := Criteria{ + ExcludeNodeIDs: []storj.NodeID{excluded}, + AutoExcludeSubnets: map[string]struct{}{}, + } + + // due to node id criteria + assert.False(t, criteria.MatchInclude(&Node{ + NodeURL: storj.NodeURL{ + ID: excluded, + Address: "192.168.0.1", + }, + })) + + // should be included as previous one excluded and + // not stored for subnet exclusion + assert.True(t, criteria.MatchInclude(&Node{ + NodeURL: storj.NodeURL{ + ID: testrand.NodeID(), + Address: "192.168.0.2", + }, + })) + +} diff --git a/satellite/nodeselection/uploadselection/node.go b/satellite/nodeselection/uploadselection/node.go index 353337cc9..68d9193bd 100644 --- a/satellite/nodeselection/uploadselection/node.go +++ b/satellite/nodeselection/uploadselection/node.go @@ -1,4 +1,4 @@ -// Copyright (C) 2020 Storj Labs, Incache. +// Copyright (C) 2020 Storj Labs, Inc. // See LICENSE for copying information. package uploadselection diff --git a/satellite/nodeselection/uploadselection/selector.go b/satellite/nodeselection/uploadselection/selector.go index 668696698..ba654a849 100644 --- a/satellite/nodeselection/uploadselection/selector.go +++ b/satellite/nodeselection/uploadselection/selector.go @@ -1,12 +1,10 @@ -// Copyright (C) 2020 Storj Labs, Incache. +// Copyright (C) 2020 Storj Labs, Inc. // 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. @@ -18,7 +16,7 @@ var _ Selector = (SelectByID)(nil) 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 { +func (nodes SelectByID) Select(n int, criteria Criteria) []*Node { if n <= 0 { return nil } @@ -27,15 +25,9 @@ func (nodes SelectByID) Select(n int, excludedIDs []storj.NodeID, excludedNets m for _, idx := range mathrand.Perm(len(nodes)) { node := nodes[idx] - if ContainsID(excludedIDs, node.ID) { + if !criteria.MatchInclude(node) { continue } - if excludedNets != nil { - if _, excluded := excludedNets[node.LastNet]; excluded { - continue - } - excludedNets[node.LastNet] = struct{}{} - } selected = append(selected, node.Clone()) if len(selected) >= n { @@ -79,7 +71,7 @@ func SelectBySubnetFromNodes(nodes []*Node) SelectBySubnet { 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 { +func (subnets SelectBySubnet) Select(n int, criteria Criteria) []*Node { if n <= 0 { return nil } @@ -89,15 +81,9 @@ func (subnets SelectBySubnet) Select(n int, excludedIDs []storj.NodeID, excluded subnet := subnets[idx] node := subnet.Nodes[mathrand.Intn(len(subnet.Nodes))] - if ContainsID(excludedIDs, node.ID) { + if !criteria.MatchInclude(node) { continue } - if excludedNets != nil { - if _, excluded := excludedNets[node.LastNet]; excluded { - continue - } - excludedNets[node.LastNet] = struct{}{} - } selected = append(selected, node.Clone()) if len(selected) >= n { @@ -107,13 +93,3 @@ func (subnets SelectBySubnet) Select(n int, excludedIDs []storj.NodeID, excluded 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 -} diff --git a/satellite/nodeselection/uploadselection/selector_test.go b/satellite/nodeselection/uploadselection/selector_test.go index 6d655313d..36ee6862e 100644 --- a/satellite/nodeselection/uploadselection/selector_test.go +++ b/satellite/nodeselection/uploadselection/selector_test.go @@ -1,4 +1,4 @@ -// Copyright (C) 2020 Storj Labs, Incache. +// Copyright (C) 2020 Storj Labs, Inc. // See LICENSE for copying information. package uploadselection_test @@ -63,7 +63,7 @@ func TestSelectByID(t *testing.T) { // perform many node selections that selects 2 nodes for i := 0; i < executionCount; i++ { - selectedNodes := selector.Select(reqCount, nil, nil) + selectedNodes := selector.Select(reqCount, uploadselection.Criteria{}) require.Len(t, selectedNodes, reqCount) for _, node := range selectedNodes { selectedNodeCount[node.ID]++ @@ -132,7 +132,7 @@ func TestSelectBySubnet(t *testing.T) { // perform many node selections that selects 2 nodes for i := 0; i < executionCount; i++ { - selectedNodes := selector.Select(reqCount, nil, map[string]struct{}{}) + selectedNodes := selector.Select(reqCount, uploadselection.Criteria{}) require.Len(t, selectedNodes, reqCount) for _, node := range selectedNodes { selectedNodeCount[node.ID]++ @@ -213,7 +213,7 @@ func TestSelectBySubnetOneAtATime(t *testing.T) { // perform many node selections that selects 1 node for i := 0; i < executionCount; i++ { - selectedNodes := selector.Select(reqCount, nil, map[string]struct{}{}) + selectedNodes := selector.Select(reqCount, uploadselection.Criteria{}) require.Len(t, selectedNodes, reqCount) for _, node := range selectedNodes { selectedNodeCount[node.ID]++ @@ -238,3 +238,58 @@ func TestSelectBySubnetOneAtATime(t *testing.T) { // expect that the single node is selected ~50% of the time assert.InDelta(t, subnetB1Count/total, uniqueSubnet, selectionEpsilon) } + +func TestSelectFiltered(t *testing.T) { + + ctx := testcontext.New(t) + defer ctx.Cleanup() + + // create 3 nodes, 2 with same subnet + lastNetDuplicate := "1.0.1" + firstID := testrand.NodeID() + subnetA1 := &uploadselection.Node{ + NodeURL: storj.NodeURL{ + ID: firstID, + Address: lastNetDuplicate + ".4:8080", + }, + LastNet: lastNetDuplicate, + LastIPPort: lastNetDuplicate + ".4:8080", + } + + secondID := testrand.NodeID() + subnetA2 := &uploadselection.Node{ + NodeURL: storj.NodeURL{ + ID: secondID, + Address: lastNetDuplicate + ".5:8080", + }, + LastNet: lastNetDuplicate, + LastIPPort: lastNetDuplicate + ".5:8080", + } + + thirdID := testrand.NodeID() + lastNetSingle := "1.0.2" + subnetB1 := &uploadselection.Node{ + NodeURL: storj.NodeURL{ + ID: thirdID, + Address: lastNetSingle + ".5:8080", + }, + LastNet: lastNetSingle, + LastIPPort: lastNetSingle + ".5:8080", + } + + nodes := []*uploadselection.Node{subnetA1, subnetA2, subnetB1} + selector := uploadselection.SelectByID(nodes) + + assert.Len(t, selector.Select(3, uploadselection.Criteria{}), 3) + assert.Len(t, selector.Select(3, uploadselection.Criteria{ExcludeNodeIDs: []storj.NodeID{firstID}}), 2) + assert.Len(t, selector.Select(3, uploadselection.Criteria{}), 3) + + assert.Len(t, selector.Select(3, uploadselection.Criteria{ExcludeNodeIDs: []storj.NodeID{firstID, secondID}}), 1) + assert.Len(t, selector.Select(3, uploadselection.Criteria{ + AutoExcludeSubnets: map[string]struct{}{}, + }), 2) + assert.Len(t, selector.Select(3, uploadselection.Criteria{ + ExcludeNodeIDs: []storj.NodeID{thirdID}, + AutoExcludeSubnets: map[string]struct{}{}, + }), 1) +} diff --git a/satellite/nodeselection/uploadselection/state.go b/satellite/nodeselection/uploadselection/state.go index 935249ce8..3cf000939 100644 --- a/satellite/nodeselection/uploadselection/state.go +++ b/satellite/nodeselection/uploadselection/state.go @@ -1,4 +1,4 @@ -// Copyright (C) 2020 Storj Labs, Incache. +// Copyright (C) 2020 Storj Labs, Inc. // See LICENSE for copying information. package uploadselection @@ -47,9 +47,9 @@ type Stats struct { 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 + // Select selects up-to n nodes which are included by the criteria. + // empty criteria includes all the nodes + Select(n int, criteria Criteria) []*Node } // NewState returns a state based on the input. @@ -100,16 +100,21 @@ func (state *State) Select(ctx context.Context, request Request) (_ []*Node, err newCount := int(float64(totalCount) * request.NewFraction) var selected []*Node - var excludedNets map[string]struct{} var reputableNodes Selector var newNodes Selector + var criteria Criteria + + if request.ExcludedIDs != nil { + criteria.ExcludeNodeIDs = request.ExcludedIDs + } + if request.Distinct { - excludedNets = map[string]struct{}{} + criteria.AutoExcludeSubnets = make(map[string]struct{}) for _, id := range request.ExcludedIDs { if net, ok := state.netByID[id]; ok { - excludedNets[net] = struct{}{} + criteria.AutoExcludeSubnets[net] = struct{}{} } } reputableNodes = state.distinct.Reputable @@ -122,12 +127,12 @@ func (state *State) Select(ctx context.Context, request Request) (_ []*Node, err // 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)...) + newNodes.Select(newCount, criteria)...) // Get all the remaining reputable nodes. reputableCount := totalCount - len(selected) selected = append(selected, - reputableNodes.Select(reputableCount, request.ExcludedIDs, excludedNets)...) + reputableNodes.Select(reputableCount, criteria)...) if len(selected) < totalCount { return selected, ErrNotEnoughNodes.New("requested from cache %d, found %d", totalCount, len(selected)) diff --git a/satellite/nodeselection/uploadselection/state_test.go b/satellite/nodeselection/uploadselection/state_test.go index 9d0d12798..539569914 100644 --- a/satellite/nodeselection/uploadselection/state_test.go +++ b/satellite/nodeselection/uploadselection/state_test.go @@ -1,4 +1,4 @@ -// Copyright (C) 2020 Storj Labs, Incache. +// Copyright (C) 2020 Storj Labs, Inc. // See LICENSE for copying information. package uploadselection_test