add default replication factor to virtual storage's configuration

Adds a configuration option for the default replication factor to
use for new repositories in a given virtual storage. This allows
Praefect to configure new repositories directly with the desired
replication factor instead of the administrator having to manually
set the replication factor of a repository after it has been created
already.

3 files changed, 77 insertions, 1 deletions

diff --git a/internal/praefect/config/config.go b/internal/praefect/config/config.go
index 6ae48474b..a82d35729 100644
--- a/internal/praefect/config/config.go
+++ b/internal/praefect/config/config.go
@@ -130,6 +130,12 @@ type Config struct {
 type VirtualStorage struct {
 	Name  string  `toml:"name"`
 	Nodes []*Node `toml:"node"`
+	// DefaultReplicationFactor is the replication factor set for new repositories.
+	// A valid value is inclusive between 1 and the number of configured storages in the
+	// virtual storage. Setting the value to 0 or below causes Praefect to not store any
+	// host assignments, falling back to the behavior of replicating to every configured
+	// storage
+	DefaultReplicationFactor int `toml:"default_replication_factor"`
 }
 
 // FromFile loads the config for the passed file path
@@ -227,6 +233,13 @@ func (c *Config) Validate() error {
 			}
 			allAddresses[node.Address] = struct{}{}
 		}
+
+		if virtualStorage.DefaultReplicationFactor > len(virtualStorage.Nodes) {
+			return fmt.Errorf(
+				"virtual storage %q has a default replication factor (%d) which is higher than the number of storages (%d)",
+				virtualStorage.Name, virtualStorage.DefaultReplicationFactor, len(virtualStorage.Nodes),
+			)
+		}
 	}
 
 	return nil
@@ -277,6 +290,17 @@ func (c *Config) StorageNames() map[string][]string {
 	return storages
 }
 
+// DefaultReplicationFactors returns a map with the default replication factors of
+// the virtual storages.
+func (c Config) DefaultReplicationFactors() map[string]int {
+	replicationFactors := make(map[string]int, len(c.VirtualStorages))
+	for _, vs := range c.VirtualStorages {
+		replicationFactors[vs.Name] = vs.DefaultReplicationFactor
+	}
+
+	return replicationFactors
+}
+
 // DB holds Postgres client configuration data.
 type DB struct {
 	Host        string `toml:"host"`
diff --git a/internal/praefect/config/config_test.go b/internal/praefect/config/config_test.go
index 9bcb23c4f..9bcfe56fb 100644
--- a/internal/praefect/config/config_test.go
+++ b/internal/praefect/config/config_test.go
@@ -180,6 +180,24 @@ func TestConfigValidation(t *testing.T) {
 			},
 			errMsg: `multiple storages have the same address`,
 		},
+		{
+			desc: "default replication factor too high",
+			changeConfig: func(cfg *Config) {
+				cfg.VirtualStorages = []*VirtualStorage{
+					{
+						Name:                     "default",
+						DefaultReplicationFactor: 2,
+						Nodes: []*Node{
+							{
+								Storage: "storage-1",
+								Address: "localhost:23456",
+							},
+						},
+					},
+				}
+			},
+			errMsg: `virtual storage "default" has a default replication factor (2) which is higher than the number of storages (1)`,
+		},
 	}
 
 	for _, tc := range testCases {
@@ -234,7 +252,8 @@ func TestConfigParsing(t *testing.T) {
 				},
 				VirtualStorages: []*VirtualStorage{
 					&VirtualStorage{
-						Name: "praefect",
+						Name:                     "praefect",
+						DefaultReplicationFactor: 2,
 						Nodes: []*Node{
 							&Node{
 								Address: "tcp://gitaly-internal-1.example.com",
@@ -425,6 +444,38 @@ func TestToPQString(t *testing.T) {
 	}
 }
 
+func TestDefaultReplicationFactors(t *testing.T) {
+	for _, tc := range []struct {
+		desc                      string
+		virtualStorages           []*VirtualStorage
+		defaultReplicationFactors map[string]int
+	}{
+		{
+			desc: "replication factors set on some",
+			virtualStorages: []*VirtualStorage{
+				{Name: "virtual-storage-1", DefaultReplicationFactor: 0},
+				{Name: "virtual-storage-2", DefaultReplicationFactor: 1},
+			},
+			defaultReplicationFactors: map[string]int{
+				"virtual-storage-1": 0,
+				"virtual-storage-2": 1,
+			},
+		},
+		{
+			desc:                      "returns always initialized map",
+			virtualStorages:           []*VirtualStorage{},
+			defaultReplicationFactors: map[string]int{},
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			require.Equal(t,
+				tc.defaultReplicationFactors,
+				Config{VirtualStorages: tc.virtualStorages}.DefaultReplicationFactors(),
+			)
+		})
+	}
+}
+
 func TestNeedsSQL(t *testing.T) {
 	testCases := []struct {
 		desc     string
diff --git a/internal/praefect/config/testdata/config.toml b/internal/praefect/config/testdata/config.toml
index 4fe025293..b402c02a5 100644
--- a/internal/praefect/config/testdata/config.toml
+++ b/internal/praefect/config/testdata/config.toml
@@ -27,6 +27,7 @@ key_path = '/home/git/key.pem'
 
 [[virtual_storage]]
 name = "praefect"
+default_replication_factor = 2
 
   [[virtual_storage.node]]
     address = "tcp://gitaly-internal-1.example.com"