[#80] netmap: process subnet in selection

Signed-off-by: Evgenii Stratonikov <evgeniy@nspcc.ru>

netmap/helper_test.go

@@ -3,6 +3,7 @@ package netmap
 import (
 	"testing"
 
+	subnetid "github.com/nspcc-dev/neofs-sdk-go/subnet/id"
 	"github.com/stretchr/testify/require"
 )
 
@@ -42,6 +43,12 @@ func newReplica(c uint32, s string) *Replica {
 	return r
 }
 
+func newSubnetID(id uint32) *subnetid.ID {
+	var s subnetid.ID
+	s.SetNumber(id)
+	return &s
+}
+
 func nodeInfoFromAttributes(props ...string) NodeInfo {
 	attrs := make([]*NodeAttribute, len(props)/2)
 	for i := range attrs {

netmap/policy.go

@@ -2,6 +2,8 @@ package netmap
 
 import (
 	"github.com/nspcc-dev/neofs-api-go/v2/netmap"
+	"github.com/nspcc-dev/neofs-api-go/v2/refs"
+	subnetid "github.com/nspcc-dev/neofs-sdk-go/subnet/id"
 )
 
 // PlacementPolicy represents v2-compatible placement policy.
@@ -32,6 +34,17 @@ func (p *PlacementPolicy) ToV2() *netmap.PlacementPolicy {
 	return (*netmap.PlacementPolicy)(p)
 }
 
+// SubnetID returns subnet to select nodes from.
+func (p *PlacementPolicy) SubnetID() *subnetid.ID {
+	return (*subnetid.ID)(
+		(*netmap.PlacementPolicy)(p).GetSubnetID())
+}
+
+// SetSubnetID sets subnet to select nodes from.
+func (p *PlacementPolicy) SetSubnetID(subnet *subnetid.ID) {
+	(*netmap.PlacementPolicy)(p).SetSubnetID((*refs.SubnetID)(subnet))
+}
+
 // Replicas returns list of object replica descriptors.
 func (p *PlacementPolicy) Replicas() []*Replica {
 	rs := (*netmap.PlacementPolicy)(p).

netmap/policy_test.go

@@ -1,13 +1,117 @@
 package netmap
 
 import (
+	"errors"
+	"strconv"
 	"testing"
 
 	"github.com/nspcc-dev/neofs-api-go/v2/netmap"
+	subnetid "github.com/nspcc-dev/neofs-sdk-go/subnet/id"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+const subnetAttrPrefix = "__NEOFS_SUBNET"
+
+func subnetAttrName(subnet uint32) string {
+	return subnetAttrPrefix + "." +
+		strconv.FormatUint(uint64(subnet), 10) + ".ENABLED"
+}
+
+func TestPlacementPolicy_Subnet(t *testing.T) {
+	nodes := []NodeInfo{
+		nodeInfoFromAttributes("ID", "0", "City", "Paris"),
+		nodeInfoFromAttributes("ID", "1", "City", "Paris"),
+		nodeInfoFromAttributes("ID", "2", "City", "London"),
+		nodeInfoFromAttributes("ID", "3", "City", "London"),
+		nodeInfoFromAttributes("ID", "4", "City", "Toronto"),
+		nodeInfoFromAttributes("ID", "5", "City", "Toronto"),
+		nodeInfoFromAttributes("ID", "6", "City", "Tokyo"),
+		nodeInfoFromAttributes("ID", "7", "City", "Tokyo"),
+	}
+	var id subnetid.ID
+	nodes[0].ExitSubnet(id)
+
+	id.SetNumber(1)
+	nodes[2].EnterSubnet(id)
+	nodes[4].EnterSubnet(id)
+
+	id.SetNumber(2)
+	nodes[5].EnterSubnet(id)
+	nodes[6].EnterSubnet(id)
+	nodes[7].EnterSubnet(id)
+
+	nm, err := NewNetmap(NodesFromInfo(nodes))
+	require.NoError(t, err)
+
+	t.Run("select 2 nodes from the default subnet in Paris", func(t *testing.T) {
+		p := newPlacementPolicy(0,
+			[]*Replica{newReplica(1, "S")},
+			[]*Selector{newSelector("S", "City", ClauseSame, 2, "F")},
+			[]*Filter{newFilter("F", "City", "Paris", OpEQ)})
+
+		_, err := nm.GetContainerNodes(p, nil)
+		require.True(t, errors.Is(err, ErrNotEnoughNodes), "got: %v", err)
+	})
+	t.Run("select 2 nodes from the default subnet in London", func(t *testing.T) {
+		p := newPlacementPolicy(0,
+			[]*Replica{newReplica(1, "S")},
+			[]*Selector{newSelector("S", "City", ClauseSame, 2, "F")},
+			[]*Filter{newFilter("F", "City", "London", OpEQ)})
+
+		v, err := nm.GetContainerNodes(p, nil)
+		require.NoError(t, err)
+
+		nodes := v.Flatten()
+		require.Equal(t, 2, len(nodes))
+		for _, n := range v.Flatten() {
+			id := n.Attribute("ID")
+			require.Contains(t, []string{"2", "3"}, id)
+		}
+	})
+	t.Run("select 2 nodes from the default subnet in Toronto", func(t *testing.T) {
+		p := newPlacementPolicy(0,
+			[]*Replica{newReplica(1, "S")},
+			[]*Selector{newSelector("S", "City", ClauseSame, 2, "F")},
+			[]*Filter{newFilter("F", "City", "Toronto", OpEQ)})
+
+		v, err := nm.GetContainerNodes(p, nil)
+		require.NoError(t, err)
+
+		nodes := v.Flatten()
+		require.Equal(t, 2, len(nodes))
+		for _, n := range v.Flatten() {
+			id := n.Attribute("ID")
+			require.Contains(t, []string{"4", "5"}, id)
+		}
+	})
+	t.Run("select 3 nodes from the non-default subnet", func(t *testing.T) {
+		p := newPlacementPolicy(0,
+			[]*Replica{newReplica(3, "")},
+			nil, nil)
+		p.SetSubnetID(newSubnetID(2))
+
+		v, err := nm.GetContainerNodes(p, nil)
+		require.NoError(t, err)
+
+		nodes := v.Flatten()
+		require.Equal(t, 3, len(nodes))
+		for _, n := range v.Flatten() {
+			id := n.Attribute("ID")
+			require.Contains(t, []string{"5", "6", "7"}, id)
+		}
+	})
+	t.Run("select nodes from the subnet via filter", func(t *testing.T) {
+		p := newPlacementPolicy(0,
+			[]*Replica{newReplica(1, "")},
+			nil,
+			[]*Filter{newFilter(MainFilterName, subnetAttrName(2), "True", OpEQ, nil)})
+
+		_, err := nm.GetContainerNodes(p, nil)
+		require.Error(t, err)
+	})
+}
+
 func TestPlacementPolicy_CBFWithEmptySelector(t *testing.T) {
 	nodes := []NodeInfo{
 		nodeInfoFromAttributes("ID", "1", "Attr", "Same"),

netmap/selector.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/nspcc-dev/hrw"
 	"github.com/nspcc-dev/neofs-api-go/v2/netmap"
+	subnetid "github.com/nspcc-dev/neofs-sdk-go/subnet/id"
 )
 
 // Selector represents v2-compatible netmap selector.
@@ -51,7 +52,7 @@ func GetNodesCount(_ *PlacementPolicy, s *Selector) (int, int) {
 // Last argument specifies if more buckets can be used to fulfill CBF.
 func (c *Context) getSelection(p *PlacementPolicy, s *Selector) ([]Nodes, error) {
 	bucketCount, nodesInBucket := GetNodesCount(p, s)
-	buckets := c.getSelectionBase(s)
+	buckets := c.getSelectionBase(p.SubnetID(), s)
 
 	if len(buckets) < bucketCount {
 		return nil, fmt.Errorf("%w: '%s'", ErrNotEnoughNodes, s.Name())
@@ -121,7 +122,7 @@ type nodeAttrPair struct {
 
 // getSelectionBase returns nodes grouped by selector attribute.
 // It it guaranteed that each pair will contain at least one node.
-func (c *Context) getSelectionBase(s *Selector) []nodeAttrPair {
+func (c *Context) getSelectionBase(subnetID *subnetid.ID, s *Selector) []nodeAttrPair {
 	f := c.Filters[s.Filter()]
 	isMain := s.Filter() == MainFilterName
 	result := []nodeAttrPair{}
@@ -129,6 +130,14 @@ func (c *Context) getSelectionBase(s *Selector) []nodeAttrPair {
 	attr := s.Attribute()
 
 	for i := range c.Netmap.Nodes {
+		var sid subnetid.ID
+		if subnetID != nil {
+			sid = *subnetID
+		}
+		// TODO(fyrchik): make `BelongsToSubnet` to accept pointer
+		if !BelongsToSubnet(c.Netmap.Nodes[i].NodeInfo, sid) {
+			continue
+		}
 		if isMain || c.match(f, c.Netmap.Nodes[i]) {
 			if attr == "" {
 				// Default attribute is transparent identifier which is different for every node.