satellite/metainfo/metabase: add node alias table

This is first step in supporting node aliases. It adds a table
that automatically assigns aliases to nodes inserted into the table.

Change-Id: Ibdf40097c3c1e5b371500203f8db203505a48adc

satellite/metainfo/metabase/alias.go

@@ -0,0 +1,79 @@
+// Copyright (C) 2021 Storj Labs, Inc.
+// See LICENSE for copying information.
+
+package metabase
+
+import (
+	"context"
+
+	"github.com/zeebo/errs"
+
+	"storj.io/common/storj"
+	"storj.io/storj/private/dbutil/pgutil"
+)
+
+// NodeAlias is a metabase local alias for NodeID-s to reduce segment table size.
+type NodeAlias int32
+
+// NodeAliasEntry is a mapping between NodeID and NodeAlias.
+type NodeAliasEntry struct {
+	ID    storj.NodeID
+	Alias NodeAlias
+}
+
+// EnsureNodeAliases contains arguments necessary for creating NodeAlias-es.
+type EnsureNodeAliases struct {
+	Nodes []storj.NodeID
+}
+
+// EnsureNodeAliases ensures that the supplied node ID-s have a alias.
+// It's safe to concurrently try and create node ID-s for the same NodeID.
+func (db *DB) EnsureNodeAliases(ctx context.Context, opts EnsureNodeAliases) (err error) {
+	defer mon.Task()(&ctx)(&err)
+
+	for _, node := range opts.Nodes {
+		if node.IsZero() {
+			return Error.New("tried to add alias to zero node")
+		}
+	}
+
+	_, err = db.db.ExecContext(ctx, `
+		INSERT INTO node_aliases(node_id)
+		SELECT unnest($1::BYTEA[])
+		ON CONFLICT DO NOTHING
+	`, pgutil.NodeIDArray(opts.Nodes))
+	if err != nil {
+		return Error.Wrap(err)
+	}
+
+	return nil
+}
+
+// ListNodeAliases lists all node alias mappings.
+func (db *DB) ListNodeAliases(ctx context.Context) (_ []NodeAliasEntry, err error) {
+	defer mon.Task()(&ctx)(&err)
+
+	var aliases []NodeAliasEntry
+	rows, err := db.db.Query(ctx, `
+		SELECT node_id, node_alias
+		FROM node_aliases
+	`)
+	if err != nil {
+		return nil, Error.New("ListNodeAliases query: %w", err)
+	}
+	defer func() { err = errs.Combine(err, rows.Close()) }()
+
+	for rows.Next() {
+		var entry NodeAliasEntry
+		err := rows.Scan(&entry.ID, &entry.Alias)
+		if err != nil {
+			return nil, Error.New("ListNodeAliases scan failed: %w", err)
+		}
+		aliases = append(aliases, entry)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, Error.New("ListNodeAliases scan failed: %w", err)
+	}
+
+	return aliases, nil
+}

satellite/metainfo/metabase/alias_test.go

@@ -0,0 +1,158 @@
+// Copyright (C) 2021 Storj Labs, Inc.
+// See LICENSE for copying information.
+
+package metabase_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/sync/errgroup"
+
+	"storj.io/common/storj"
+	"storj.io/common/testcontext"
+	"storj.io/common/testrand"
+	"storj.io/storj/satellite/metainfo/metabase"
+)
+
+func TestNodeAliases(t *testing.T) {
+	All(t, func(ctx *testcontext.Context, t *testing.T, db *metabase.DB) {
+		t.Run("Zero", func(t *testing.T) {
+			defer DeleteAll{}.Check(ctx, t, db)
+
+			nodes := []storj.NodeID{
+				testrand.NodeID(),
+				{},
+			}
+			EnsureNodeAliases{
+				Opts: metabase.EnsureNodeAliases{
+					Nodes: nodes,
+				},
+				ErrClass: &metabase.Error,
+				ErrText:  "tried to add alias to zero node",
+			}.Check(ctx, t, db)
+		})
+
+		t.Run("Empty", func(t *testing.T) {
+			defer DeleteAll{}.Check(ctx, t, db)
+
+			aliasesAfter := ListNodeAliases{}.Check(ctx, t, db)
+			require.Len(t, aliasesAfter, 0)
+		})
+
+		t.Run("Valid", func(t *testing.T) {
+			defer DeleteAll{}.Check(ctx, t, db)
+
+			nodes := []storj.NodeID{
+				testrand.NodeID(),
+				testrand.NodeID(),
+				testrand.NodeID(),
+			}
+			nodes = append(nodes, nodes...) // add duplicates to our slice
+
+			EnsureNodeAliases{
+				Opts: metabase.EnsureNodeAliases{
+					Nodes: nodes,
+				},
+			}.Check(ctx, t, db)
+
+			EnsureNodeAliases{
+				Opts: metabase.EnsureNodeAliases{
+					Nodes: nodes,
+				},
+			}.Check(ctx, t, db)
+
+			aliases := ListNodeAliases{}.Check(ctx, t, db)
+			require.Len(t, aliases, 3)
+
+			for _, entry := range aliases {
+				require.True(t, nodesContains(nodes, entry.ID))
+				require.LessOrEqual(t, int(entry.Alias), 3)
+			}
+
+			EnsureNodeAliases{
+				Opts: metabase.EnsureNodeAliases{
+					Nodes: []storj.NodeID{testrand.NodeID()},
+				},
+			}.Check(ctx, t, db)
+
+			aliasesAfter := ListNodeAliases{}.Check(ctx, t, db)
+			require.Len(t, aliasesAfter, 4)
+		})
+
+		t.Run("Concurrent", func(t *testing.T) {
+			defer DeleteAll{}.Check(ctx, t, db)
+
+			nodes := make([]storj.NodeID, 128)
+			for i := range nodes {
+				nodes[i] = testrand.NodeID()
+			}
+
+			var group errgroup.Group
+			for k := range nodes {
+				node := nodes[k]
+				group.Go(func() error {
+					return db.EnsureNodeAliases(ctx, metabase.EnsureNodeAliases{
+						Nodes: []storj.NodeID{node},
+					})
+				})
+			}
+			require.NoError(t, group.Wait())
+
+			aliases := ListNodeAliases{}.Check(ctx, t, db)
+			seen := map[metabase.NodeAlias]bool{}
+			require.Len(t, aliases, len(nodes))
+			for _, entry := range aliases {
+				require.True(t, nodesContains(nodes, entry.ID))
+				require.LessOrEqual(t, int(entry.Alias), len(nodes))
+
+				require.False(t, seen[entry.Alias])
+				seen[entry.Alias] = true
+			}
+		})
+
+		t.Run("Stress Concurrent", func(t *testing.T) {
+			defer DeleteAll{}.Check(ctx, t, db)
+
+			nodes := make([]storj.NodeID, 128)
+			for i := range nodes {
+				nodes[i] = testrand.NodeID()
+			}
+			group, gctx := errgroup.WithContext(ctx)
+			for k := 0; k < 16; k++ {
+				group.Go(func() error {
+					loc := nodes
+					for len(loc) > 0 {
+						k := testrand.Intn(10)
+						if k > len(loc) {
+							k = len(loc)
+						}
+						var batch []storj.NodeID
+						batch, loc = loc[:k], loc[k:]
+						err := db.EnsureNodeAliases(gctx,
+							metabase.EnsureNodeAliases{Nodes: batch},
+						)
+						if err != nil {
+							panic(err)
+						}
+
+						if gctx.Err() != nil {
+							break
+						}
+					}
+					return nil
+				})
+			}
+			require.NoError(t, group.Wait())
+		})
+	})
+}
+
+func nodesContains(nodes []storj.NodeID, v storj.NodeID) bool {
+	for _, n := range nodes {
+		if n == v {
+			return true
+		}
+	}
+	return false
+}

satellite/metainfo/metabase/db.go

@@ -60,6 +60,8 @@ func (db *DB) DestroyTables(ctx context.Context) error {
 	_, err := db.db.ExecContext(ctx, `
 		DROP TABLE IF EXISTS objects;
 		DROP TABLE IF EXISTS segments;
+		DROP TABLE IF EXISTS node_aliases;
+		DROP SEQUENCE IF EXISTS node_alias_seq;
 	`)
 	return Error.Wrap(err)
 }
@@ -140,6 +142,23 @@ func (db *DB) PostgresMigration() *migrate.Migration {
 					`ALTER TABLE objects ALTER COLUMN total_encrypted_size TYPE INT8;`,
 				},
 			},
+			{
+				DB:          &db.db,
+				Description: "add node aliases table",
+				Version:     3,
+				Action: migrate.SQL{
+					// We use a custom sequence to ensure small alias values.
+					`CREATE SEQUENCE node_alias_seq
+						INCREMENT BY 1
+						MINVALUE 1 MAXVALUE 2147483647 -- MaxInt32
+						START WITH 1
+					`,
+					`CREATE TABLE node_aliases (
+						node_id    BYTEA  NOT NULL UNIQUE,
+						node_alias INT4   NOT NULL UNIQUE default nextval('node_alias_seq')
+					)`,
+				},
+			},
 		},
 	}
 }

satellite/metainfo/metabase/raw.go

@@ -87,6 +87,8 @@ func (db *DB) TestingDeleteAll(ctx context.Context) (err error) {
 	_, err = db.db.ExecContext(ctx, `
 		DELETE FROM objects;
 		DELETE FROM segments;
+		DELETE FROM node_aliases;
+		SELECT setval('node_alias_seq', 1, false);
 	`)
 	return Error.Wrap(err)
 }

satellite/metainfo/metabase/test_test.go

@@ -433,6 +433,29 @@ func (step IterateObjectsWithStatus) Check(ctx *testcontext.Context, t testing.T
 	require.Zero(t, diff)
 }
 
+type EnsureNodeAliases struct {
+	Opts metabase.EnsureNodeAliases
+
+	ErrClass *errs.Class
+	ErrText  string
+}
+
+func (step EnsureNodeAliases) Check(ctx *testcontext.Context, t testing.TB, db *metabase.DB) {
+	err := db.EnsureNodeAliases(ctx, step.Opts)
+	checkError(t, err, step.ErrClass, step.ErrText)
+}
+
+type ListNodeAliases struct {
+	ErrClass *errs.Class
+	ErrText  string
+}
+
+func (step ListNodeAliases) Check(ctx *testcontext.Context, t testing.TB, db *metabase.DB) []metabase.NodeAliasEntry {
+	result, err := db.ListNodeAliases(ctx)
+	checkError(t, err, step.ErrClass, step.ErrText)
+	return result
+}
+
 func checkError(t testing.TB, err error, errClass *errs.Class, errText string) {
 	if errClass != nil {
 		require.True(t, errClass.Has(err), "expected an error %v got %v", *errClass, err)