diff --git a/netmap/attributes.go b/netmap/attributes.go index 07ce2686..e1be340a 100644 --- a/netmap/attributes.go +++ b/netmap/attributes.go @@ -113,12 +113,16 @@ func WriteSubnetInfo(node *NodeInfo, info NodeSubnetInfo) { // ErrRemoveSubnet is returned when a node needs to leave the subnet. var ErrRemoveSubnet = errors.New("remove subnet") +var errNoSubnets = errors.New("no subnets") + // IterateSubnets iterates over all subnets the node belongs to and passes the IDs to f. // Handler must not be nil. // -// If f returns ErrRemoveSubnet, then removes subnet entry. Breaks on any other non-nil error and returns it. +// If f returns ErrRemoveSubnet, then removes subnet entry. Note that this leads to an instant mutation of NodeInfo. +// Breaks on any other non-nil error and returns it. // // Returns an error if any subnet attribute has wrong format. +// Returns an error if the node is not included in any subnet by the end of the loop. func IterateSubnets(node *NodeInfo, f func(refs.SubnetID) error) error { attrs := node.GetAttributes() @@ -126,6 +130,7 @@ func IterateSubnets(node *NodeInfo, f func(refs.SubnetID) error) error { err error id refs.SubnetID metZero bool // if zero subnet's attribute was met in for-loop + entries uint ) for i := 0; i < len(attrs); i++ { // range must not be used because of attrs mutation in body @@ -168,6 +173,7 @@ func IterateSubnets(node *NodeInfo, f func(refs.SubnetID) error) error { if !isRemoveErr { // no handler's error and non-zero subnet + entries++ continue } else if metZero { // removal error and zero subnet. @@ -183,7 +189,10 @@ func IterateSubnets(node *NodeInfo, f func(refs.SubnetID) error) error { // we can set False or remove attribute, latter is more memory/network efficient. attrs = append(attrs[:i], attrs[i+1:]...) i-- + continue } + + entries++ } if !metZero { @@ -199,9 +208,15 @@ func IterateSubnets(node *NodeInfo, f func(refs.SubnetID) error) error { attr.SetValue(attrSubnetValExit) attrs = append(attrs, &attr) + } else { + entries++ } } + if entries <= 0 { + return errNoSubnets + } + node.SetAttributes(attrs) return nil diff --git a/netmap/attributes_test.go b/netmap/attributes_test.go index 5cee47ee..7e03b8cf 100644 --- a/netmap/attributes_test.go +++ b/netmap/attributes_test.go @@ -179,6 +179,15 @@ func TestSubnets(t *testing.T) { t.Run("zero", func(t *testing.T) { var node netmap.NodeInfo + // enter to some non-zero subnet so that zero is not the only one + var attr netmap.Attribute + + attr.SetKey(subnetAttrKey("321")) + attr.SetValue("True") + + attrs := []*netmap.Attribute{&attr} + node.SetAttributes(attrs) + err := netmap.IterateSubnets(&node, func(id refs.SubnetID) error { if refs.IsZeroSubnet(&id) { return netmap.ErrRemoveSubnet @@ -189,12 +198,19 @@ func TestSubnets(t *testing.T) { require.NoError(t, err) - attrs := node.GetAttributes() - require.Len(t, attrs, 1) + attrs = node.GetAttributes() + require.Len(t, attrs, 2) - attr := attrs[0] - assertSubnetAttrKey(t, attr, 0) - require.Equal(t, "False", attr.GetValue()) + found := false + + for i := range attrs { + if attrs[i].GetKey() == subnetAttrKey("0") { + require.Equal(t, "False", attrs[i].GetValue()) + found = true + } + } + + require.True(t, found) }) t.Run("non-zero", func(t *testing.T) { @@ -222,5 +238,30 @@ func TestSubnets(t *testing.T) { attrs = node.GetAttributes() require.Empty(t, attrs) }) + + t.Run("all", func(t *testing.T) { + var ( + node netmap.NodeInfo + attrs []*netmap.Attribute + ) + + // enter to some non-zero subnet so that zero is not the only one + for i := 1; i <= 5; i++ { + var attr netmap.Attribute + + attr.SetKey(subnetAttrKey(strconv.Itoa(i))) + attr.SetValue("True") + + attrs = append(attrs, &attr) + } + + node.SetAttributes(attrs) + + err := netmap.IterateSubnets(&node, func(id refs.SubnetID) error { + return netmap.ErrRemoveSubnet + }) + + require.Error(t, err) + }) }) }