implement functionality to set a repository's replication factor

This commit adds SetReplicationFactor method on the AssignmentStore
for assigning and unassigning storage nodes from a repository until
the desired replication factor is met.

Repository's assignments define where the repository should be
replicated to. Praefect already queries the assignments in the
relevant locations. If no assignments are set, Praefect replicates
the repository to every storage node. This can be wasteful as the
repository's replication factor increases with the number of storages
in the virtual storage. SetReplicationFactor allows for setting a
replication factor for a repository to support replicating to only
a specified number of storage nodes instead.

2 files changed, 265 insertions, 0 deletions

diff --git a/internal/praefect/datastore/assignment.go b/internal/praefect/datastore/assignment.go
index 6bab0a594..25018693e 100644
--- a/internal/praefect/datastore/assignment.go
+++ b/internal/praefect/datastore/assignment.go
@@ -12,6 +12,18 @@ func newVirtualStorageNotFoundError(virtualStorage string) error {
 	return fmt.Errorf("virtual storage %q not found", virtualStorage)
 }
 
+func newUnattainableReplicationFactorError(attempted, maximum int) error {
+	return fmt.Errorf("attempted to set replication factor %d but virtual storage only contains %d storages", attempted, maximum)
+}
+
+func newMinimumReplicationFactorError(replicationFactor int) error {
+	return fmt.Errorf("attempted to set replication factor %d but minimum is 1", replicationFactor)
+}
+
+func newRepositoryNotFoundError(virtualStorage, relativePath string) error {
+	return fmt.Errorf("repository %q/%q not found", virtualStorage, relativePath)
+}
+
 // AssignmentStore manages host assignments in Postgres.
 type AssignmentStore struct {
 	db                 glsql.Querier
@@ -61,3 +73,126 @@ AND   storage = ANY($3)
 
 	return assignedStorages, nil
 }
+
+// SetReplicationFactor assigns or unassigns host nodes from the repository to meet the desired replication factor.
+// SetReplicationFactor returns an error when trying to set a replication factor that exceeds the storage node count
+// in the virtual storage. An error is also returned when trying to set a replication factor below one. The primary node
+// won't be unassigned as it needs a copy of the repository to accept writes. Likewise, the primary is the first storage
+// that gets assigned when setting a replication factor for a repository. Assignments of unconfigured storages are ignored.
+// This might cause the actual replication factor to be higher than desired if the replication factor is set during an upgrade
+// from a Praefect node that does not yet know about a new node. As assignments of unconfigured storages are ignored, replication
+// factor of repositories assigned to a storage node removed from the cluster is effectively decreased.
+func (s AssignmentStore) SetReplicationFactor(ctx context.Context, virtualStorage, relativePath string, replicationFactor int) ([]string, error) {
+	candidateStorages, ok := s.configuredStorages[virtualStorage]
+	if !ok {
+		return nil, fmt.Errorf("unknown virtual storage: %q", virtualStorage)
+	}
+
+	if replicationFactor < 1 {
+		return nil, newMinimumReplicationFactorError(replicationFactor)
+	}
+
+	if max := len(candidateStorages); replicationFactor > max {
+		return nil, newUnattainableReplicationFactorError(replicationFactor, max)
+	}
+
+	// The query works as follows:
+	//
+	// 1. `repository` CTE locks the repository's record for the duration of the update.
+	//    This prevents concurrent updates to the `repository_assignments` table for the given
+	//    repository. It is not sufficient to rely on row locks in `repository_assignments`
+	//    as there might be rows being inserted or deleted in another transaction that
+	//    our transaction does not lock. This could be the case if the replication factor
+	//    is being increased concurrently from two different nodes and they assign different
+	//    storages.
+	//
+	// 2. `existing_assignments` CTE gets the existing assignments for the repository. While
+	//    there may be assignments in the database for storage nodes that were removed from the
+	//    cluster, the query filters them out.
+	//
+	// 3. `created_assignments` CTE assigns new hosts to the repository if the replication
+	//    factor has been increased. Random storages which are not yet assigned to the repository
+	//    are picked until the replication factor is met. The primary of a repository is always
+	//    assigned first.
+	//
+	// 4. `removed_assignments` CTE removes host assignments if the replication factor has been
+	//    decreased. Primary is never removed as it needs a copy of the repository in order to
+	//    accept writes. Random hosts are removed until the replication factor is met.
+	//
+	// 6. Finally we return the current set of assignments. CTE updates are not visible in the
+	//    tables during the transaction. To account for that, we filter out removed assignments
+	//    from the existing assignments. If the replication factor was increased, we'll include the
+	//    created assignments. If the replication factor did not change, the query returns the
+	//    current assignments.
+	rows, err := s.db.QueryContext(ctx, `
+WITH repository AS (
+	SELECT virtual_storage, relative_path, "primary"
+	FROM repositories
+	WHERE virtual_storage = $1
+	AND   relative_path   = $2
+	FOR UPDATE
+),
+
+existing_assignments AS (
+	SELECT storage
+	FROM repository
+	JOIN repository_assignments USING (virtual_storage, relative_path)
+	WHERE storage = ANY($4::text[])
+),
+
+created_assignments AS (
+	INSERT INTO repository_assignments
+	SELECT virtual_storage, relative_path, storage
+	FROM repository
+	CROSS JOIN ( SELECT unnest($4::text[]) AS storage ) AS configured_storages
+	WHERE storage NOT IN ( SELECT storage FROM existing_assignments )
+	ORDER BY CASE WHEN storage = "primary" THEN 1 ELSE 0 END DESC, random()
+	LIMIT ( SELECT GREATEST(COUNT(*), $3) - COUNT(*) FROM existing_assignments )
+	RETURNING storage
+),
+
+removed_assignments AS (
+	DELETE FROM repository_assignments
+	USING (
+		SELECT virtual_storage, relative_path, storage
+		FROM repository, existing_assignments
+		WHERE storage != "primary"
+		ORDER BY random()
+		LIMIT ( SELECT COUNT(*) - LEAST(COUNT(*), $3)  FROM existing_assignments )
+	) AS removals
+	WHERE repository_assignments.virtual_storage = removals.virtual_storage
+	AND   repository_assignments.relative_path   = removals.relative_path
+	AND   repository_assignments.storage         = removals.storage
+	RETURNING removals.storage
+)
+
+SELECT storage
+FROM existing_assignments
+WHERE storage NOT IN ( SELECT storage FROM removed_assignments )
+UNION
+SELECT storage
+FROM created_assignments
+ORDER BY storage
+	`, virtualStorage, relativePath, replicationFactor, pq.StringArray(candidateStorages))
+	if err != nil {
+		return nil, fmt.Errorf("query: %w", err)
+	}
+
+	defer rows.Close()
+
+	var storages []string
+	for rows.Next() {
+		var storage string
+		if err := rows.Scan(&storage); err != nil {
+			return nil, fmt.Errorf("scan: %w", err)
+		}
+
+		storages = append(storages, storage)
+	}
+
+	if len(storages) == 0 {
+		return nil, newRepositoryNotFoundError(virtualStorage, relativePath)
+	}
+
+	return storages, rows.Err()
+}
diff --git a/internal/praefect/datastore/assignment_test.go b/internal/praefect/datastore/assignment_test.go
index 5baeda66d..96f9add51 100644
--- a/internal/praefect/datastore/assignment_test.go
+++ b/internal/praefect/datastore/assignment_test.go
@@ -98,3 +98,133 @@ func TestAssignmentStore_GetHostAssignments(t *testing.T) {
 		})
 	}
 }
+
+func TestAssignmentStore_SetReplicationFactor(t *testing.T) {
+	type matcher func(testing.TB, []string)
+
+	equal := func(expected []string) matcher {
+		return func(t testing.TB, actual []string) {
+			t.Helper()
+			require.Equal(t, expected, actual)
+		}
+	}
+
+	contains := func(expecteds ...[]string) matcher {
+		return func(t testing.TB, actual []string) {
+			t.Helper()
+			require.Contains(t, expecteds, actual)
+		}
+	}
+
+	for _, tc := range []struct {
+		desc                  string
+		existingAssignments   []string
+		nonExistentRepository bool
+		replicationFactor     int
+		requireStorages       matcher
+		error                 error
+	}{
+		{
+			desc:                  "increase replication factor of non-existent repository",
+			nonExistentRepository: true,
+			replicationFactor:     1,
+			error:                 newRepositoryNotFoundError("virtual-storage", "relative-path"),
+		},
+		{
+			desc:              "primary prioritized when setting the first assignments",
+			replicationFactor: 1,
+			requireStorages:   equal([]string{"primary"}),
+		},
+		{
+			desc:                "increasing replication factor ignores unconfigured storages",
+			existingAssignments: []string{"unconfigured-storage"},
+			replicationFactor:   1,
+			requireStorages:     equal([]string{"primary"}),
+		},
+		{
+			desc:                "replication factor already achieved",
+			existingAssignments: []string{"primary", "secondary-1"},
+			replicationFactor:   2,
+			requireStorages:     equal([]string{"primary", "secondary-1"}),
+		},
+		{
+			desc:                "increase replication factor by a step",
+			existingAssignments: []string{"primary"},
+			replicationFactor:   2,
+			requireStorages:     contains([]string{"primary", "secondary-1"}, []string{"primary", "secondary-2"}),
+		},
+		{
+			desc:                "increase replication factor to maximum",
+			existingAssignments: []string{"primary"},
+			replicationFactor:   3,
+			requireStorages:     equal([]string{"primary", "secondary-1", "secondary-2"}),
+		},
+		{
+			desc:                "increased replication factor unattainable",
+			existingAssignments: []string{"primary"},
+			replicationFactor:   4,
+			error:               newUnattainableReplicationFactorError(4, 3),
+		},
+		{
+			desc:                "decreasing replication factor ignores unconfigured storages",
+			existingAssignments: []string{"secondary-1", "unconfigured-storage"},
+			replicationFactor:   1,
+			requireStorages:     equal([]string{"secondary-1"}),
+		},
+		{
+			desc:                "decrease replication factor by a step",
+			existingAssignments: []string{"primary", "secondary-1", "secondary-2"},
+			replicationFactor:   2,
+			requireStorages:     contains([]string{"primary", "secondary-1"}, []string{"primary", "secondary-2"}),
+		},
+		{
+			desc:                "decrease replication factor to minimum",
+			existingAssignments: []string{"primary", "secondary-1", "secondary-2"},
+			replicationFactor:   1,
+			requireStorages:     equal([]string{"primary"}),
+		},
+		{
+			desc:              "minimum replication factor is enforced",
+			replicationFactor: 0,
+			error:             newMinimumReplicationFactorError(0),
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			ctx, cancel := testhelper.Context()
+			defer cancel()
+
+			db := getDB(t)
+
+			configuredStorages := map[string][]string{"virtual-storage": {"primary", "secondary-1", "secondary-2"}}
+
+			if !tc.nonExistentRepository {
+				_, err := db.ExecContext(ctx, `
+					INSERT INTO repositories (virtual_storage, relative_path, "primary")
+					VALUES ('virtual-storage', 'relative-path', 'primary')
+				`)
+				require.NoError(t, err)
+			}
+
+			for _, storage := range tc.existingAssignments {
+				_, err := db.ExecContext(ctx, `
+					INSERT INTO repository_assignments VALUES ('virtual-storage', 'relative-path', $1)
+				`, storage)
+				require.NoError(t, err)
+			}
+
+			store := NewAssignmentStore(db, configuredStorages)
+
+			setStorages, err := store.SetReplicationFactor(ctx, "virtual-storage", "relative-path", tc.replicationFactor)
+			require.Equal(t, tc.error, err)
+			if tc.error != nil {
+				return
+			}
+
+			tc.requireStorages(t, setStorages)
+
+			assignedStorages, err := store.GetHostAssignments(ctx, "virtual-storage", "relative-path")
+			require.NoError(t, err)
+			tc.requireStorages(t, assignedStorages)
+		})
+	}
+}