support storing assignments on repository creation

When variable replication factor is enabled, the repository's host
nodes should be selected when the repository is being created. The
selected assignments should then be persisted so they can be used
for future routing decisions. This commit adds support for storing
assignments in RepositoryStore's CreateRepository.

The functionality is behind a boolean switch as assignments are not
supported when repository specific primaries are not used. Separate
switch from storing primary is used in order to support repository
specific primaries without variable replication factor. This allows
for independently rolling out the features to production later.

3 files changed, 88 insertions, 23 deletions

diff --git a/internal/praefect/datastore/repository_store.go b/internal/praefect/datastore/repository_store.go
index 2cf36be9a..bc2c8c8ba 100644
--- a/internal/praefect/datastore/repository_store.go
+++ b/internal/praefect/datastore/repository_store.go
@@ -103,7 +103,10 @@ type RepositoryStore interface {
 	//
 	// storePrimary should be set when repository specific primaries are enabled. When set, the primary is stored as
 	// the repository's primary.
-	CreateRepository(ctx context.Context, virtualStorage, relativePath, primary string, storePrimary bool) error
+	//
+	// storeAssignments should be set when variable replication factor is enabled. When set, the primary and the
+	// secondaries are stored as the assigned hosts of the repository.
+	CreateRepository(ctx context.Context, virtualStorage, relativePath, primary string, secondaries []string, storePrimary, storeAssignments bool) error
 	// DeleteRepository deletes the repository from the virtual storage and the storage. Returns
 	// RepositoryNotExistsError when trying to delete a repository which has no record in the virtual storage
 	// or the storage.
@@ -301,7 +304,7 @@ AND storage = ANY($3)
 
 //nolint:stylecheck
 //nolint:golint
-func (rs *PostgresRepositoryStore) CreateRepository(ctx context.Context, virtualStorage, relativePath, primary string, storePrimary bool) error {
+func (rs *PostgresRepositoryStore) CreateRepository(ctx context.Context, virtualStorage, relativePath, primary string, secondaries []string, storePrimary, storeAssignments bool) error {
 	const q = `
 WITH repo AS (
 	INSERT INTO repositories (
@@ -310,7 +313,23 @@ WITH repo AS (
 		generation,
 		"primary"
 	) VALUES ($1, $2, 0, CASE WHEN $4 THEN $3 END)
+),
+
+assignments AS (
+	INSERT INTO repository_assignments (
+		virtual_storage,
+		relative_path,
+		storage
+	)
+	SELECT $1, $2, storage
+	FROM (
+		SELECT unnest($5::text[]) AS storage
+		UNION
+		SELECT $3
+	) AS storages
+	WHERE $6
 )
+
 INSERT INTO storage_repositories (
 	virtual_storage,
 	relative_path,
@@ -320,7 +339,8 @@ INSERT INTO storage_repositories (
 VALUES ($1, $2, $3, 0)
 `
 
-	_, err := rs.db.ExecContext(ctx, q, virtualStorage, relativePath, primary, storePrimary)
+	_, err := rs.db.ExecContext(ctx, q,
+		virtualStorage, relativePath, primary, storePrimary, pq.StringArray(secondaries), storeAssignments)
 
 	var pqerr *pq.Error
 	if errors.As(err, &pqerr) && pqerr.Code.Name() == "unique_violation" {
diff --git a/internal/praefect/datastore/repository_store_mock.go b/internal/praefect/datastore/repository_store_mock.go
index ea051d1d0..8df265ef5 100644
--- a/internal/praefect/datastore/repository_store_mock.go
+++ b/internal/praefect/datastore/repository_store_mock.go
@@ -9,7 +9,7 @@ type MockRepositoryStore struct {
 	IncrementGenerationFunc                func(ctx context.Context, virtualStorage, relativePath, primary string, secondaries []string) error
 	GetReplicatedGenerationFunc            func(ctx context.Context, virtualStorage, relativePath, source, target string) (int, error)
 	SetGenerationFunc                      func(ctx context.Context, virtualStorage, relativePath, storage string, generation int) error
-	CreateRepositoryFunc                   func(ctx context.Context, virtualStorage, relativePath, storage string, storePrimary bool) error
+	CreateRepositoryFunc                   func(ctx context.Context, virtualStorage, relativePath, primary string, secondaries []string, storePrimary, storeAssignments bool) error
 	DeleteRepositoryFunc                   func(ctx context.Context, virtualStorage, relativePath, storage string) error
 	RenameRepositoryFunc                   func(ctx context.Context, virtualStorage, relativePath, storage, newRelativePath string) error
 	GetConsistentStoragesFunc              func(ctx context.Context, virtualStorage, relativePath string) (map[string]struct{}, error)
@@ -52,12 +52,12 @@ func (m MockRepositoryStore) SetGeneration(ctx context.Context, virtualStorage,
 
 //nolint:stylecheck
 //nolint:golint
-func (m MockRepositoryStore) CreateRepository(ctx context.Context, virtualStorage, relativePath, storage string, storePrimary bool) error {
+func (m MockRepositoryStore) CreateRepository(ctx context.Context, virtualStorage, relativePath, primary string, secondaries []string, storePrimary, storeAssignments bool) error {
 	if m.CreateRepositoryFunc == nil {
 		return nil
 	}
 
-	return m.CreateRepositoryFunc(ctx, virtualStorage, relativePath, storage, storePrimary)
+	return m.CreateRepositoryFunc(ctx, virtualStorage, relativePath, primary, secondaries, storePrimary, storeAssignments)
 }
 
 func (m MockRepositoryStore) DeleteRepository(ctx context.Context, virtualStorage, relativePath, storage string) error {
diff --git a/internal/praefect/datastore/repository_store_test.go b/internal/praefect/datastore/repository_store_test.go
index efcfbe41f..365260a03 100644
--- a/internal/praefect/datastore/repository_store_test.go
+++ b/internal/praefect/datastore/repository_store_test.go
@@ -7,13 +7,15 @@ import (
 	"database/sql"
 	"testing"
 
+	"github.com/lib/pq"
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
 )
 
 // repositoryRecord represents Praefect's records related to a repository.
 type repositoryRecord struct {
-	primary string
+	primary     string
+	assignments []string
 }
 
 // virtualStorageStates represents the virtual storage's view of which repositories should exist.
@@ -34,8 +36,14 @@ func TestRepositoryStore_Postgres(t *testing.T) {
 
 		requireVirtualStorageState := func(t *testing.T, ctx context.Context, exp virtualStorageState) {
 			rows, err := db.QueryContext(ctx, `
-SELECT virtual_storage, relative_path, "primary"
+SELECT virtual_storage, relative_path, "primary", assigned_storages
 FROM repositories
+LEFT JOIN (
+	SELECT virtual_storage, relative_path, array_agg(storage ORDER BY storage) AS assigned_storages
+	FROM repository_assignments
+	GROUP BY virtual_storage, relative_path
+) AS repository_assignments USING (virtual_storage, relative_path)
+
 				`)
 			require.NoError(t, err)
 			defer rows.Close()
@@ -45,14 +53,16 @@ FROM repositories
 				var (
 					virtualStorage, relativePath string
 					primary                      sql.NullString
+					assignments                  pq.StringArray
 				)
-				require.NoError(t, rows.Scan(&virtualStorage, &relativePath, &primary))
+				require.NoError(t, rows.Scan(&virtualStorage, &relativePath, &primary, &assignments))
 				if act[virtualStorage] == nil {
 					act[virtualStorage] = make(map[string]repositoryRecord)
 				}
 
 				act[virtualStorage][relativePath] = repositoryRecord{
-					primary: primary.String,
+					primary:     primary.String,
+					assignments: assignments,
 				}
 			}
 
@@ -299,34 +309,69 @@ func testRepositoryStore(t *testing.T, newStore repositoryStoreFactory) {
 	t.Run("CreateRepository", func(t *testing.T) {
 		t.Run("successfully created", func(t *testing.T) {
 			for _, tc := range []struct {
-				desc         string
-				storePrimary bool
-				primary      string
+				desc                string
+				secondaries         []string
+				storePrimary        bool
+				storeAssignments    bool
+				expectedPrimary     string
+				expectedAssignments []string
 			}{
 				{
-					desc:         "primary not stored",
-					storePrimary: false,
+					desc: "store only repository record",
+				},
+				{
+					desc:            "primary stored",
+					secondaries:     []string{"secondary-1", "secondary-2"},
+					storePrimary:    true,
+					expectedPrimary: "primary",
+				},
+				{
+					desc:                "assignments stored",
+					storeAssignments:    true,
+					secondaries:         []string{"secondary-1", "secondary-2"},
+					expectedAssignments: []string{"primary", "secondary-1", "secondary-2"},
 				},
 				{
-					desc:         "primary stored",
-					storePrimary: true,
-					primary:      stor,
+					desc:                "store primary and assignments",
+					storePrimary:        true,
+					storeAssignments:    true,
+					secondaries:         []string{"secondary-1", "secondary-2"},
+					expectedPrimary:     "primary",
+					expectedAssignments: []string{"primary", "secondary-1", "secondary-2"},
+				},
+				{
+					desc:                "store primary and no secondaries",
+					storePrimary:        true,
+					storeAssignments:    true,
+					secondaries:         []string{},
+					expectedPrimary:     "primary",
+					expectedAssignments: []string{"primary"},
+				},
+				{
+					desc:                "store primary and nil secondaries",
+					storePrimary:        true,
+					storeAssignments:    true,
+					expectedPrimary:     "primary",
+					expectedAssignments: []string{"primary"},
 				},
 			} {
 				t.Run(tc.desc, func(t *testing.T) {
 					rs, requireState := newStore(t, nil)
 
-					require.NoError(t, rs.CreateRepository(ctx, vs, repo, stor, tc.storePrimary))
+					require.NoError(t, rs.CreateRepository(ctx, vs, repo, "primary", tc.secondaries, tc.storePrimary, tc.storeAssignments))
 					requireState(t, ctx,
 						virtualStorageState{
 							vs: {
-								repo: repositoryRecord{primary: tc.primary},
+								repo: repositoryRecord{
+									primary:     tc.expectedPrimary,
+									assignments: tc.expectedAssignments,
+								},
 							},
 						},
 						storageState{
 							vs: {
 								repo: {
-									stor: 0,
+									"primary": 0,
 								},
 							},
 						},
@@ -338,10 +383,10 @@ func testRepositoryStore(t *testing.T, newStore repositoryStoreFactory) {
 		t.Run("conflict", func(t *testing.T) {
 			rs, _ := newStore(t, nil)
 
-			require.NoError(t, rs.CreateRepository(ctx, vs, repo, stor, false))
+			require.NoError(t, rs.CreateRepository(ctx, vs, repo, stor, nil, false, false))
 			require.Equal(t,
 				RepositoryExistsError{vs, repo, stor},
-				rs.CreateRepository(ctx, vs, repo, stor, false),
+				rs.CreateRepository(ctx, vs, repo, stor, nil, false, false),
 			)
 		})
 	})