support default replication factor in PerRepositoryRouter

This commit adds support for selecting host nodes randomly on
repository creation in PerRepositoryRouter. New repositories get
random secondaries assigned as host nodes until the default
replication factor is met. If the virtual storage has no default
replication factor configured, every configured node, except the
randomly picked primary,  is included as secondaries. This is
fallback behavior to match Praefect's behavior prior to variable
replication factor.

4 files changed, 126 insertions, 48 deletions

diff --git a/cmd/praefect/main.go b/cmd/praefect/main.go
index a4692647c..b59dc560c 100644
--- a/cmd/praefect/main.go
+++ b/cmd/praefect/main.go
@@ -338,6 +338,7 @@ func run(cfgs []starter.Config, conf config.Config) error {
 			praefect.NewLockedRandom(rand.New(rand.NewSource(time.Now().UnixNano()))),
 			rs,
 			assignmentStore,
+			map[string]int{},
 		)
 	} else {
 		healthChecker = praefect.HealthChecker(nodeManager)
diff --git a/internal/praefect/coordinator_test.go b/internal/praefect/coordinator_test.go
index c6bbb37ca..7cba80328 100644
--- a/internal/praefect/coordinator_test.go
+++ b/internal/praefect/coordinator_test.go
@@ -699,6 +699,7 @@ func TestStreamDirector_repo_creation(t *testing.T) {
 					},
 					nil,
 					nil,
+					map[string]int{},
 				)
 			default:
 				t.Fatalf("unexpected election strategy: %q", tc.electionStrategy)
diff --git a/internal/praefect/router_per_repository.go b/internal/praefect/router_per_repository.go
index 5e95f006d..930170254 100644
--- a/internal/praefect/router_per_repository.go
+++ b/internal/praefect/router_per_repository.go
@@ -52,23 +52,25 @@ type PrimaryGetter interface {
 
 // PerRepositoryRouter implements a router that routes requests respecting per repository primary nodes.
 type PerRepositoryRouter struct {
-	conns Connections
-	ag    AssignmentGetter
-	pg    PrimaryGetter
-	rand  Random
-	hc    HealthChecker
-	rs    datastore.RepositoryStore
+	conns                     Connections
+	ag                        AssignmentGetter
+	pg                        PrimaryGetter
+	rand                      Random
+	hc                        HealthChecker
+	rs                        datastore.RepositoryStore
+	defaultReplicationFactors map[string]int
 }
 
 // NewPerRepositoryRouter returns a new PerRepositoryRouter using the passed configuration.
-func NewPerRepositoryRouter(conns Connections, pg PrimaryGetter, hc HealthChecker, rand Random, rs datastore.RepositoryStore, ag AssignmentGetter) *PerRepositoryRouter {
+func NewPerRepositoryRouter(conns Connections, pg PrimaryGetter, hc HealthChecker, rand Random, rs datastore.RepositoryStore, ag AssignmentGetter, defaultReplicationFactors map[string]int) *PerRepositoryRouter {
 	return &PerRepositoryRouter{
-		conns: conns,
-		pg:    pg,
-		rand:  rand,
-		hc:    hc,
-		rs:    rs,
-		ag:    ag,
+		conns:                     conns,
+		pg:                        pg,
+		rand:                      rand,
+		hc:                        hc,
+		rs:                        rs,
+		ag:                        ag,
+		defaultReplicationFactors: defaultReplicationFactors,
 	}
 }
 
@@ -225,6 +227,11 @@ func (r *PerRepositoryRouter) RouteRepositoryCreation(ctx context.Context, virtu
 		return RepositoryMutatorRoute{}, err
 	}
 
+	replicationFactor := r.defaultReplicationFactors[virtualStorage]
+	if replicationFactor == 1 {
+		return RepositoryMutatorRoute{Primary: primary}, nil
+	}
+
 	// NodeManagerRouter doesn't consider any secondaries as consistent when creating a repository,
 	// thus the primary is the only participant in the transaction and the secondaries get replicated to.
 	// PerRepositoryRouter matches that behavior here for consistency.
@@ -237,6 +244,17 @@ func (r *PerRepositoryRouter) RouteRepositoryCreation(ctx context.Context, virtu
 		replicationTargets = append(replicationTargets, storage)
 	}
 
+	// replicationFactor being zero indicates it has not been configured. If so, we fallback to the behavior
+	// of no assignments and replicate everywhere.
+	if replicationFactor > 1 {
+		r.rand.Shuffle(len(replicationTargets), func(i, j int) {
+			replicationTargets[i], replicationTargets[j] = replicationTargets[j], replicationTargets[i]
+		})
+
+		// deduct one as the primary is also hosting the repository
+		replicationTargets = replicationTargets[:replicationFactor-1]
+	}
+
 	return RepositoryMutatorRoute{
 		Primary:            primary,
 		ReplicationTargets: replicationTargets,
diff --git a/internal/praefect/router_per_repository_test.go b/internal/praefect/router_per_repository_test.go
index aa5553f23..ff4c1e83c 100644
--- a/internal/praefect/router_per_repository_test.go
+++ b/internal/praefect/router_per_repository_test.go
@@ -94,6 +94,7 @@ func TestPerRepositoryRouter_RouteStorageAccessor(t *testing.T) {
 				},
 				nil,
 				nil,
+				nil,
 			)
 
 			node, err := router.RouteStorageAccessor(ctx, tc.virtualStorage)
@@ -202,6 +203,7 @@ func TestPerRepositoryRouter_RouteRepositoryAccessor(t *testing.T) {
 					},
 				},
 				nil,
+				nil,
 			)
 
 			node, err := router.RouteRepositoryAccessor(ctx, tc.virtualStorage, "repository")
@@ -342,6 +344,7 @@ func TestPerRepositoryRouter_RouteRepositoryMutator(t *testing.T) {
 					},
 				},
 				tc.assignedNodes,
+				nil,
 			)
 
 			route, err := router.RouteRepositoryMutator(ctx, tc.virtualStorage, "repository")
@@ -373,44 +376,93 @@ func TestPerRepositoryRouter_RouteRepositoryCreation(t *testing.T) {
 		"virtual-storage-1": {"primary", "secondary-1", "secondary-2"},
 	}
 
+	type matcher func(*testing.T, []string)
+
+	requireOneOf := func(expected ...[]string) func(*testing.T, []string) {
+		return func(t *testing.T, actual []string) {
+			sort.Strings(actual)
+			require.Contains(t, expected, actual)
+		}
+	}
+
+	requireNil := func(t *testing.T, actual []string) {
+		require.Nil(t, actual)
+	}
+
 	for _, tc := range []struct {
-		desc               string
-		virtualStorage     string
-		healthyNodes       StaticHealthChecker
-		numCandidates      int
-		pickCandidate      int
-		primary            string
-		replicationTargets []string
-		error              error
+		desc                      string
+		virtualStorage            string
+		healthyNodes              StaticHealthChecker
+		replicationFactor         int
+		primaryCandidates         int
+		primaryPick               int
+		secondaryCandidates       int
+		primary                   string
+		requireReplicationTargets matcher
+		error                     error
 	}{
 		{
-			desc:           "no healthy nodes",
-			virtualStorage: "virtual-storage-1",
-			healthyNodes:   StaticHealthChecker{},
-			error:          ErrNoHealthyNodes,
+			desc:                      "no healthy nodes",
+			virtualStorage:            "virtual-storage-1",
+			healthyNodes:              StaticHealthChecker{},
+			error:                     ErrNoHealthyNodes,
+			requireReplicationTargets: requireNil,
 		},
 		{
-			desc:           "invalid virtual storage",
-			virtualStorage: "invalid",
-			error:          nodes.ErrVirtualStorageNotExist,
+			desc:                      "invalid virtual storage",
+			virtualStorage:            "invalid",
+			error:                     nodes.ErrVirtualStorageNotExist,
+			requireReplicationTargets: requireNil,
 		},
 		{
-			desc:               "no healthy secondaries",
-			virtualStorage:     "virtual-storage-1",
-			healthyNodes:       StaticHealthChecker{"virtual-storage-1": {"primary"}},
-			numCandidates:      1,
-			pickCandidate:      0,
-			primary:            "primary",
-			replicationTargets: []string{"secondary-1", "secondary-2"},
+			desc:                      "no healthy secondaries",
+			virtualStorage:            "virtual-storage-1",
+			healthyNodes:              StaticHealthChecker{"virtual-storage-1": {"primary"}},
+			primaryCandidates:         1,
+			primaryPick:               0,
+			primary:                   "primary",
+			requireReplicationTargets: requireOneOf([]string{"secondary-1", "secondary-2"}),
 		},
 		{
-			desc:               "success",
-			virtualStorage:     "virtual-storage-1",
-			healthyNodes:       StaticHealthChecker(configuredNodes),
-			numCandidates:      3,
-			pickCandidate:      0,
-			primary:            "primary",
-			replicationTargets: []string{"secondary-1", "secondary-2"},
+			desc:                      "success",
+			virtualStorage:            "virtual-storage-1",
+			healthyNodes:              StaticHealthChecker(configuredNodes),
+			primaryCandidates:         3,
+			primaryPick:               0,
+			primary:                   "primary",
+			requireReplicationTargets: requireOneOf([]string{"secondary-1", "secondary-2"}),
+		},
+		{
+			desc:                      "replication factor of one configured",
+			virtualStorage:            "virtual-storage-1",
+			healthyNodes:              StaticHealthChecker(configuredNodes),
+			replicationFactor:         1,
+			primaryCandidates:         3,
+			primaryPick:               0,
+			primary:                   "primary",
+			requireReplicationTargets: requireNil,
+		},
+		{
+			desc:                      "replication factor of two configured",
+			virtualStorage:            "virtual-storage-1",
+			healthyNodes:              StaticHealthChecker(configuredNodes),
+			replicationFactor:         2,
+			primaryCandidates:         3,
+			primaryPick:               0,
+			secondaryCandidates:       2,
+			primary:                   "primary",
+			requireReplicationTargets: requireOneOf([]string{"secondary-1"}, []string{"secondary-2"}),
+		},
+		{
+			desc:                      "replication factor of three configured with unhealthy secondary",
+			virtualStorage:            "virtual-storage-1",
+			healthyNodes:              StaticHealthChecker{"virtual-storage-1": {"primary", "secondary-1"}},
+			replicationFactor:         3,
+			primaryCandidates:         2,
+			primaryPick:               0,
+			secondaryCandidates:       2,
+			primary:                   "primary",
+			requireReplicationTargets: requireOneOf([]string{"secondary-1", "secondary-2"}),
 		},
 	} {
 		t.Run(tc.desc, func(t *testing.T) {
@@ -431,20 +483,26 @@ func TestPerRepositoryRouter_RouteRepositoryCreation(t *testing.T) {
 				tc.healthyNodes,
 				mockRandom{
 					intnFunc: func(n int) int {
-						require.Equal(t, tc.numCandidates, n)
-						return tc.pickCandidate
+						require.Equal(t, tc.primaryCandidates, n)
+						return tc.primaryPick
+					},
+					shuffleFunc: func(n int, swap func(i, j int)) {
+						require.Equal(t, tc.secondaryCandidates, n)
 					},
 				},
 				nil,
 				nil,
+				map[string]int{"virtual-storage-1": tc.replicationFactor},
 			).RouteRepositoryCreation(ctx, tc.virtualStorage)
+			require.Equal(t, tc.error, err)
 
-			sort.Strings(route.ReplicationTargets)
+			// assert replication targets separately as the picked secondary
+			// is random
+			tc.requireReplicationTargets(t, route.ReplicationTargets)
+			route.ReplicationTargets = nil
 
-			require.Equal(t, tc.error, err)
 			require.Equal(t, RepositoryMutatorRoute{
-				Primary:            RouterNode{Storage: tc.primary, Connection: conns[tc.virtualStorage][tc.primary]},
-				ReplicationTargets: tc.replicationTargets,
+				Primary: RouterNode{Storage: tc.primary, Connection: conns[tc.virtualStorage][tc.primary]},
 			}, route)
 		})
 	}