satellite/nodeselection: improve annotation composability

We would like to make it easier to accept multiple annotations.

Examples:
```
country("GB") && annotation(...)
annotated(annotated(X,...),...)
```

Change-Id: I92e622e8b985b314dadddf83b17976c245eb2069
This commit is contained in:
Márton Elek 2023-08-23 15:47:49 +02:00 committed by Storj Robot
parent 2cdc1a973f
commit 5c12a3406d
4 changed files with 123 additions and 32 deletions

View File

@ -15,10 +15,50 @@ type NodeFilter interface {
MatchInclude(node *SelectedNode) bool
}
// NodeFilterWithAnnotation is a NodeFilter with additional annotations.
type NodeFilterWithAnnotation interface {
NodeFilter
GetAnnotation(name string) string
}
// Annotation can be used as node filters in 'XX && annotation('...')' like struct.
type Annotation struct {
Key string
Value string
}
// MatchInclude implements NodeFilter.
func (a Annotation) MatchInclude(node *SelectedNode) bool {
return true
}
// GetAnnotation implements NodeFilterWithAnnotation.
func (a Annotation) GetAnnotation(name string) string {
if a.Key == name {
return a.Value
}
return ""
}
var _ NodeFilterWithAnnotation = Annotation{}
// AnnotatedNodeFilter is just a NodeFilter with additional annotations.
type AnnotatedNodeFilter struct {
Filter NodeFilter
Annotations map[string]string
Annotations []Annotation
}
// GetAnnotation implements NodeFilterWithAnnotation.
func (a AnnotatedNodeFilter) GetAnnotation(name string) string {
for _, a := range a.Annotations {
if a.Key == name {
return a.Value
}
}
if annotated, ok := a.Filter.(NodeFilterWithAnnotation); ok {
return annotated.GetAnnotation(name)
}
return ""
}
// MatchInclude implements NodeFilter.
@ -27,35 +67,27 @@ func (a AnnotatedNodeFilter) MatchInclude(node *SelectedNode) bool {
}
// WithAnnotation adds annotations to a NodeFilter.
func WithAnnotation(filter NodeFilter, name string, value string) NodeFilter {
if anf, ok := filter.(AnnotatedNodeFilter); ok {
anf.Annotations[name] = value
return anf
}
func WithAnnotation(filter NodeFilter, name string, value string) NodeFilterWithAnnotation {
return AnnotatedNodeFilter{
Filter: filter,
Annotations: map[string]string{
name: value,
Annotations: []Annotation{
{
Key: name,
Value: value,
},
},
}
}
// GetAnnotation retrieves annotation from AnnotatedNodeFilter.
func GetAnnotation(filter NodeFilter, name string) string {
if annotated, ok := filter.(AnnotatedNodeFilter); ok {
return annotated.Annotations[name]
}
if filters, ok := filter.(NodeFilters); ok {
for _, filter := range filters {
if annotated, ok := filter.(AnnotatedNodeFilter); ok {
return annotated.Annotations[name]
}
}
if annotated, ok := filter.(NodeFilterWithAnnotation); ok {
return annotated.GetAnnotation(name)
}
return ""
}
var _ NodeFilter = AnnotatedNodeFilter{}
var _ NodeFilterWithAnnotation = AnnotatedNodeFilter{}
// NodeFilters is a collection of multiple node filters (all should vote with true).
type NodeFilters []NodeFilter
@ -94,7 +126,20 @@ func (n NodeFilters) WithExcludedIDs(ds []storj.NodeID) NodeFilters {
return append(n, ExcludedIDs(ds))
}
var _ NodeFilter = NodeFilters{}
// GetAnnotation implements NodeFilterWithAnnotation.
func (n NodeFilters) GetAnnotation(name string) string {
for _, filter := range n {
if annotated, ok := filter.(NodeFilterWithAnnotation); ok {
value := annotated.GetAnnotation(name)
if value != "" {
return value
}
}
}
return ""
}
var _ NodeFilterWithAnnotation = NodeFilters{}
// CountryFilter can select nodes based on the condition of the country code.
type CountryFilter struct {

View File

@ -55,6 +55,20 @@ func TestCriteria_ExcludedNodeNetworks(t *testing.T) {
}))
}
func TestAnnotations(t *testing.T) {
k := WithAnnotation(NodeFilters{}, "foo", "bar")
require.Equal(t, "bar", k.GetAnnotation("foo"))
k = NodeFilters{WithAnnotation(NodeFilters{}, "foo", "bar")}
require.Equal(t, "bar", k.GetAnnotation("foo"))
k = Annotation{
Key: "foo",
Value: "bar",
}
require.Equal(t, "bar", k.GetAnnotation("foo"))
}
func TestCriteria_Geofencing(t *testing.T) {
eu := NodeFilters{}.WithCountryFilter(location.NewSet(EuCountries...))
us := NodeFilters{}.WithCountryFilter(location.NewSet(location.UnitedStates))

View File

@ -151,15 +151,16 @@ func (d *ConfigurablePlacementRule) AddPlacementFromString(definitions string) e
}
return res, nil
},
"annotated": func(filter nodeselection.NodeFilter, kv map[string]string) (nodeselection.AnnotatedNodeFilter, error) {
"annotated": func(filter nodeselection.NodeFilter, kv ...nodeselection.Annotation) (nodeselection.AnnotatedNodeFilter, error) {
return nodeselection.AnnotatedNodeFilter{
Filter: filter,
Annotations: kv,
}, nil
},
"annotation": func(key string, value string) (map[string]string, error) {
return map[string]string{
key: value,
"annotation": func(key string, value string) (nodeselection.Annotation, error) {
return nodeselection.Annotation{
Key: key,
Value: value,
}, nil
},
"exclude": func(filter nodeselection.NodeFilter) (nodeselection.NodeFilter, error) {

View File

@ -166,16 +166,47 @@ func TestPlacementFromString(t *testing.T) {
}))
})
t.Run("annotated", func(t *testing.T) {
p := NewPlacementRules()
err := p.AddPlacementFromString(`11:annotated(country("GB"),annotation("autoExcludeSubnet","off"))`)
require.NoError(t, err)
filters := p.placements[storj.PlacementConstraint(11)]
require.True(t, filters.MatchInclude(&nodeselection.SelectedNode{
CountryCode: location.UnitedKingdom,
}))
t.Run("annotation usage", func(t *testing.T) {
t.Run("normal", func(t *testing.T) {
t.Parallel()
p := NewPlacementRules()
err := p.AddPlacementFromString(`11:annotated(country("GB"),annotation("autoExcludeSubnet","off"))`)
require.NoError(t, err)
filters := p.placements[storj.PlacementConstraint(11)]
require.True(t, filters.MatchInclude(&nodeselection.SelectedNode{
CountryCode: location.UnitedKingdom,
}))
require.Equal(t, nodeselection.GetAnnotation(filters, "autoExcludeSubnet"), "off")
require.Equal(t, nodeselection.GetAnnotation(filters, "autoExcludeSubnet"), "off")
})
t.Run("with &&", func(t *testing.T) {
t.Parallel()
p := NewPlacementRules()
err := p.AddPlacementFromString(`11:country("GB") && annotation("foo","bar") && annotation("bar","foo")`)
require.NoError(t, err)
filters := p.placements[storj.PlacementConstraint(11)]
require.True(t, filters.MatchInclude(&nodeselection.SelectedNode{
CountryCode: location.UnitedKingdom,
}))
require.Equal(t, "bar", nodeselection.GetAnnotation(filters, "foo"))
require.Equal(t, "foo", nodeselection.GetAnnotation(filters, "bar"))
require.Equal(t, "", nodeselection.GetAnnotation(filters, "kossuth"))
})
t.Run("chained", func(t *testing.T) {
t.Parallel()
p := NewPlacementRules()
err := p.AddPlacementFromString(`11:annotated(annotated(country("GB"),annotation("foo","bar")),annotation("bar","foo"))`)
require.NoError(t, err)
filters := p.placements[storj.PlacementConstraint(11)]
require.True(t, filters.MatchInclude(&nodeselection.SelectedNode{
CountryCode: location.UnitedKingdom,
}))
require.Equal(t, "bar", nodeselection.GetAnnotation(filters, "foo"))
require.Equal(t, "foo", nodeselection.GetAnnotation(filters, "bar"))
require.Equal(t, "", nodeselection.GetAnnotation(filters, "kossuth"))
})
})
t.Run("exclude", func(t *testing.T) {