From 5c12a3406d053ce91d1b68ea05954f3a9de78c04 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=C3=A1rton=20Elek?= Date: Wed, 23 Aug 2023 15:47:49 +0200 Subject: [PATCH] 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 --- satellite/nodeselection/filter.go | 83 ++++++++++++++++++++------ satellite/nodeselection/filter_test.go | 14 +++++ satellite/overlay/placement.go | 9 +-- satellite/overlay/placement_test.go | 49 ++++++++++++--- 4 files changed, 123 insertions(+), 32 deletions(-) diff --git a/satellite/nodeselection/filter.go b/satellite/nodeselection/filter.go index c6c092e7b..100a9d1f6 100644 --- a/satellite/nodeselection/filter.go +++ b/satellite/nodeselection/filter.go @@ -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 { diff --git a/satellite/nodeselection/filter_test.go b/satellite/nodeselection/filter_test.go index fa871e4aa..50d3186ff 100644 --- a/satellite/nodeselection/filter_test.go +++ b/satellite/nodeselection/filter_test.go @@ -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)) diff --git a/satellite/overlay/placement.go b/satellite/overlay/placement.go index 3bbbaba9d..679344ee3 100644 --- a/satellite/overlay/placement.go +++ b/satellite/overlay/placement.go @@ -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) { diff --git a/satellite/overlay/placement_test.go b/satellite/overlay/placement_test.go index e7968d416..53ed27b36 100644 --- a/satellite/overlay/placement_test.go +++ b/satellite/overlay/placement_test.go @@ -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) {