satellite/{nodeselection,overlay}: use location.Set

location.Set is faster for comparisons.

Updates #6028

Change-Id: I764eb5cafc507f908e4168b16a7994cc7721ce4d

satellite/nodeselection/filter.go

@@ -43,10 +43,8 @@ func (n NodeFilters) MatchInclude(node *SelectedNode) bool {
 }
 
 // WithCountryFilter is a helper to create a new filter with additional CountryFilter.
-func (n NodeFilters) WithCountryFilter(filter func(code location.CountryCode) bool) NodeFilters {
+func (n NodeFilters) WithCountryFilter(permit location.Set) NodeFilters {
-	return append(n, CountryFilter{
+	return append(n, NewCountryFilter(permit))
-		matchIncludeCountry: filter,
-	})
 }
 
 // WithAutoExcludeSubnets is a helper to create a new filter with additional AutoExcludeSubnets.
@@ -81,25 +79,27 @@ var _ NodeFilter = CountryCodeExclude{}
 
 // CountryFilter can select nodes based on the condition of the country code.
 type CountryFilter struct {
-	matchIncludeCountry func(code location.CountryCode) bool
+	permit location.Set
 }
 
 // NewCountryFilter creates a new CountryFilter.
-func NewCountryFilter(filter func(code location.CountryCode) bool) NodeFilter {
+func NewCountryFilter(permit location.Set) NodeFilter {
-	return CountryFilter{
+	return &CountryFilter{
-		matchIncludeCountry: filter,
+		permit: permit,
 	}
 }
 
 // MatchInclude implements NodeFilter interface.
-func (p CountryFilter) MatchInclude(node *SelectedNode) bool {
+func (p *CountryFilter) MatchInclude(node *SelectedNode) bool {
-	return p.matchIncludeCountry(node.CountryCode)
+	return p.permit.Contains(node.CountryCode)
 }
 
-var _ NodeFilter = CountryFilter{}
+var _ NodeFilter = &CountryFilter{}
 
 // AutoExcludeSubnets pick at most one node from network.
+//
 // Stateful!!! should be re-created for each new selection request.
+// It should only be used as the last filter.
 type AutoExcludeSubnets struct {
 	seenSubnets map[string]struct{}
 }

satellite/nodeselection/filter_test.go

@@ -73,18 +73,8 @@ func TestCriteria_NodeIDAndSubnet(t *testing.T) {
 }
 
 func TestCriteria_Geofencing(t *testing.T) {
-	eu := NodeFilters{}.WithCountryFilter(func(code location.CountryCode) bool {
+	eu := NodeFilters{}.WithCountryFilter(EuCountries)
-		for _, c := range EuCountries {
+	us := NodeFilters{}.WithCountryFilter(location.NewSet(location.UnitedStates))
-			if c == code {
-				return true
-			}
-		}
-		return false
-	})
-
-	us := NodeFilters{}.WithCountryFilter(func(code location.CountryCode) bool {
-		return code == location.UnitedStates
-	})
 
 	cases := []struct {
 		name     string

satellite/nodeselection/region.go

@@ -6,7 +6,7 @@ package nodeselection
 import "storj.io/common/storj/location"
 
 // EuCountries defines the member countries of European Union.
-var EuCountries = []location.CountryCode{
+var EuCountries = location.NewSet(
 	location.Austria,
 	location.Belgium,
 	location.Bulgaria,
@@ -34,10 +34,10 @@ var EuCountries = []location.CountryCode{
 	location.Slovakia,
 	location.Spain,
 	location.Sweden,
-}
+)
 
 // EeaCountries defined the EEA countries.
-var EeaCountries = append(EuCountries,
+var EeaCountries = EuCountries.With(
 	location.Iceland,
 	location.Liechtenstein,
 	location.Norway,

satellite/overlay/placement.go

@@ -64,31 +64,11 @@ func NewPlacementRules() *ConfigurablePlacementRule {
 
 // AddLegacyStaticRules initializes all the placement rules defined earlier in static golang code.
 func (d *ConfigurablePlacementRule) AddLegacyStaticRules() {
-	d.placements[storj.EEA] = nodeselection.NodeFilters{}.WithCountryFilter(func(isoCountryCode location.CountryCode) bool {
+	d.placements[storj.EEA] = nodeselection.NodeFilters{nodeselection.NewCountryFilter(nodeselection.EeaCountries)}
-		for _, c := range nodeselection.EeaCountries {
+	d.placements[storj.EU] = nodeselection.NodeFilters{nodeselection.NewCountryFilter(nodeselection.EuCountries)}
-			if c == isoCountryCode {
+	d.placements[storj.US] = nodeselection.NodeFilters{nodeselection.NewCountryFilter(location.NewSet(location.UnitedStates))}
-				return true
+	d.placements[storj.DE] = nodeselection.NodeFilters{nodeselection.NewCountryFilter(location.NewSet(location.Germany))}
-			}
+	d.placements[storj.NR] = nodeselection.NodeFilters{nodeselection.NewCountryFilter(location.NewFullSet().Without(location.Russia, location.Belarus))}
-		}
-		return false
-	})
-	d.placements[storj.EU] = nodeselection.NodeFilters{}.WithCountryFilter(func(isoCountryCode location.CountryCode) bool {
-		for _, c := range nodeselection.EuCountries {
-			if c == isoCountryCode {
-				return true
-			}
-		}
-		return false
-	})
-	d.placements[storj.US] = nodeselection.NodeFilters{}.WithCountryFilter(func(isoCountryCode location.CountryCode) bool {
-		return isoCountryCode == location.UnitedStates
-	})
-	d.placements[storj.DE] = nodeselection.NodeFilters{}.WithCountryFilter(func(isoCountryCode location.CountryCode) bool {
-		return isoCountryCode == location.Germany
-	})
-	d.placements[storj.NR] = nodeselection.NodeFilters{}.WithCountryFilter(func(isoCountryCode location.CountryCode) bool {
-		return isoCountryCode != location.Russia && isoCountryCode != location.Belarus
-	})
 }
 
 // AddPlacementRule registers a new placement.
@@ -100,22 +80,15 @@ func (d *ConfigurablePlacementRule) AddPlacementRule(id storj.PlacementConstrain
 func (d *ConfigurablePlacementRule) AddPlacementFromString(definitions string) error {
 	env := map[any]any{
 		"country": func(countries ...string) (nodeselection.NodeFilters, error) {
-			countryCodes := make([]location.CountryCode, len(countries))
+			var set location.Set
-			for i, country := range countries {
+			for _, country := range countries {
 				code := location.ToCountryCode(country)
 				if code == location.None {
 					return nil, errs.New("invalid country code %q", code)
 				}
-				countryCodes[i] = code
+				set.Include(code)
 			}
-			return nodeselection.NodeFilters{}.WithCountryFilter(func(code location.CountryCode) bool {
+			return nodeselection.NodeFilters{nodeselection.NewCountryFilter(set)}, nil
-				for _, expectedCode := range countryCodes {
-					if code == expectedCode {
-						return true
-					}
-				}
-				return false
-			}), nil
 		},
 		"all": func(filters ...nodeselection.NodeFilters) (nodeselection.NodeFilters, error) {
 			res := nodeselection.NodeFilters{}
@@ -166,7 +139,7 @@ func (d *ConfigurablePlacementRule) AddPlacementFromString(definitions string) e
 
 // CreateFilters implements PlacementCondition.
 func (d *ConfigurablePlacementRule) CreateFilters(constraint storj.PlacementConstraint) (filter nodeselection.NodeFilters) {
-	if constraint == 0 {
+	if constraint == storj.EveryCountry {
 		return nodeselection.NodeFilters{}
 	}
 	if filters, found := d.placements[constraint]; found {

satellite/overlay/service.go

@@ -325,14 +325,15 @@ func NewService(log *zap.Logger, db DB, nodeEvents nodeevents.DB, placementRules
 	defaultSelection := nodeselection.NodeFilters{}
 
 	if len(config.Node.UploadExcludedCountryCodes) > 0 {
-		defaultSelection = defaultSelection.WithCountryFilter(func(code location.CountryCode) bool {
+		set := location.NewFullSet()
-			for _, nodeCountry := range config.Node.UploadExcludedCountryCodes {
+		for _, country := range config.Node.UploadExcludedCountryCodes {
-				if nodeCountry == code.String() {
+			countryCode := location.ToCountryCode(country)
-					return false
+			if countryCode == location.None {
+				return nil, Error.New("invalid country %q", country)
 			}
+			set.Remove(countryCode)
 		}
-			return true
+		defaultSelection = defaultSelection.WithCountryFilter(set)
-		})
 	}
 
 	uploadSelectionCache, err := NewUploadSelectionCache(log, db,