satellite/overlay: Add ability to exclude country codes on upload

Create global config to specify a list of country codes that should be
excluded from node selection during uploads.

This exclusion is not implemented when the upload selection cache is
disabled.

Change-Id: Ic41e8b4f18857a11045668eac23107da99668a72

satellite/nodeselection/uploadselection/criteria.go

@@ -5,13 +5,15 @@ package uploadselection
 
 import (
 	"storj.io/common/storj"
+	"storj.io/common/storj/location"
 )
 
 // 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.
-	Placement          storj.PlacementConstraint
+	ExcludeNodeIDs       []storj.NodeID
+	AutoExcludeSubnets   map[string]struct{} // initialize it with empty map to keep only one node per subnet.
+	Placement            storj.PlacementConstraint
+	ExcludedCountryCodes []location.CountryCode
 }
 
 // MatchInclude returns with true if node is selected.
@@ -30,6 +32,13 @@ func (c *Criteria) MatchInclude(node *Node) bool {
 		}
 		c.AutoExcludeSubnets[node.LastNet] = struct{}{}
 	}
+
+	for _, code := range c.ExcludedCountryCodes {
+		if node.CountryCode == code {
+			return false
+		}
+	}
+
 	return true
 }
 

satellite/nodeselection/uploadselection/state.go

@@ -10,6 +10,7 @@ import (
 	"github.com/zeebo/errs"
 
 	"storj.io/common/storj"
+	"storj.io/common/storj/location"
 )
 
 // ErrNotEnoughNodes is when selecting nodes failed with the given parameters.
@@ -83,11 +84,12 @@ func NewState(reputableNodes, newNodes []*Node) *State {
 
 // Request contains arguments for State.Request.
 type Request struct {
-	Count       int
-	NewFraction float64
-	Distinct    bool
-	ExcludedIDs []storj.NodeID
-	Placement   storj.PlacementConstraint
+	Count                int
+	NewFraction          float64
+	Distinct             bool
+	ExcludedIDs          []storj.NodeID
+	Placement            storj.PlacementConstraint
+	ExcludedCountryCodes []string
 }
 
 // Select selects requestedCount nodes where there will be newFraction nodes.
@@ -111,6 +113,10 @@ func (state *State) Select(ctx context.Context, request Request) (_ []*Node, err
 		criteria.ExcludeNodeIDs = request.ExcludedIDs
 	}
 
+	for _, code := range request.ExcludedCountryCodes {
+		criteria.ExcludedCountryCodes = append(criteria.ExcludedCountryCodes, location.ToCountryCode(code))
+	}
+
 	criteria.Placement = request.Placement
 
 	if request.Distinct {

satellite/overlay/config.go

@@ -44,6 +44,8 @@ type NodeSelectionConfig struct {
 	MinimumDiskSpace memory.Size   `help:"how much disk space a node at minimum must have to be selected for upload" default:"500.00MB" testDefault:"100.00MB"`
 
 	AsOfSystemTime AsOfSystemTimeConfig
+
+	UploadExcludedCountryCodes []string `help:"list of country codes to exclude from node selection for uploads" default:"" testDefault:"FR,BE"`
 }
 
 // GeoIPConfig is a configuration struct that helps configure the GeoIP lookup features on the satellite.

satellite/overlay/service.go

@@ -398,6 +398,7 @@ func (service *Service) FindStorageNodesForUpload(ctx context.Context, req FindS
 		req.AsOfSystemInterval = service.config.Node.AsOfSystemTime.DefaultInterval
 	}
 
+	// TODO excluding country codes on upload if cache is disabled is not implemented
 	if service.config.NodeSelectionCache.Disabled {
 		return service.FindStorageNodesWithPreferences(ctx, req, &service.config.Node)
 	}

satellite/overlay/uploadselection.go

@@ -106,11 +106,12 @@ func (cache *UploadSelectionCache) GetNodes(ctx context.Context, req FindStorage
 	}
 
 	selected, err := state.Select(ctx, uploadselection.Request{
-		Count:       req.RequestedCount,
-		NewFraction: cache.selectionConfig.NewNodeFraction,
-		Distinct:    cache.selectionConfig.DistinctIP,
-		ExcludedIDs: req.ExcludedIDs,
-		Placement:   req.Placement,
+		Count:                req.RequestedCount,
+		NewFraction:          cache.selectionConfig.NewNodeFraction,
+		Distinct:             cache.selectionConfig.DistinctIP,
+		ExcludedIDs:          req.ExcludedIDs,
+		Placement:            req.Placement,
+		ExcludedCountryCodes: cache.selectionConfig.UploadExcludedCountryCodes,
 	})
 	if uploadselection.ErrNotEnoughNodes.Has(err) {
 		err = ErrNotEnoughNodes.Wrap(err)

satellite/overlay/uploadselection_test.go

@@ -20,6 +20,7 @@ import (
 	"storj.io/common/sync2"
 	"storj.io/common/testcontext"
 	"storj.io/common/testrand"
+	"storj.io/storj/private/testplanet"
 	"storj.io/storj/satellite"
 	"storj.io/storj/satellite/overlay"
 	"storj.io/storj/satellite/satellitedb/satellitedbtest"
@@ -226,6 +227,28 @@ func TestGetNodes(t *testing.T) {
 	})
 }
 
+func TestGetNodesExcludeCountryCodes(t *testing.T) {
+	testplanet.Run(t, testplanet.Config{
+		SatelliteCount: 1, StorageNodeCount: 2, UplinkCount: 0,
+	}, func(t *testing.T, ctx *testcontext.Context, planet *testplanet.Planet) {
+		err := planet.Satellites[0].Overlay.Service.TestNodeCountryCode(ctx, planet.StorageNodes[0].ID(), "FR")
+		require.NoError(t, err)
+
+		cache := planet.Satellites[0].Overlay.Service.UploadSelectionCache
+
+		// confirm cache.GetNodes returns the correct nodes
+		selectedNodes, err := cache.GetNodes(ctx, overlay.FindStorageNodesRequest{RequestedCount: 2})
+		// we only expect one node to be returned, even though we requested two, so there will be an error
+		require.Error(t, err)
+
+		_, new := cache.Size()
+		require.Equal(t, 2, new)
+		require.Equal(t, 1, len(selectedNodes))
+		// the node that was returned should be the one that does not have the "FR" country code
+		require.Equal(t, planet.StorageNodes[1].ID(), selectedNodes[0].ID)
+	})
+}
+
 func TestGetNodesConcurrent(t *testing.T) {
 	ctx := testcontext.New(t)
 	defer ctx.Cleanup()

scripts/testdata/satellite-config.yaml.lock

@@ -613,6 +613,9 @@ identity.key-path: /root/.local/share/storj/identity/satellite/identity.key
 # the amount of time without seeing a node before its considered offline
 # overlay.node.online-window: 4h0m0s
 
+# list of country codes to exclude from node selection for uploads
+# overlay.node.upload-excluded-country-codes: '[]'
+
 # list of country codes to exclude nodes from target repair selection
 # overlay.repair-excluded-country-codes: '[]'