Use repository generations to determine the best leader to elect

We should take into account repositories that are not yet exist
on the storage, so we should treat them as completely out of date.
With this we won't have situation when storage will be chosen as
primary because it has only one repository that is up to date, but
missing all other repositories that exist on other storages.

Closes: https://gitlab.com/gitlab-org/gitaly/-/issues/3008

2 files changed, 49 insertions, 11 deletions

diff --git a/internal/praefect/nodes/sql_elector.go b/internal/praefect/nodes/sql_elector.go
index c6e728496..f5dbd1c9e 100644
--- a/internal/praefect/nodes/sql_elector.go
+++ b/internal/praefect/nodes/sql_elector.go
@@ -348,13 +348,14 @@ func (s *sqlElector) electNewPrimary(candidates []*sqlCandidate) error {
 	}
 
 	q := `
-		SELECT sr.storage
+		SELECT storages.storage
 		FROM repositories AS r
-		JOIN storage_repositories AS sr USING(virtual_storage, relative_path)
-		WHERE sr.storage = ANY($1)
-			AND r.virtual_storage = $2
-		GROUP BY sr.storage
-		ORDER BY SUM(r.generation - sr.generation)`
+		CROSS JOIN (SELECT UNNEST($1::TEXT[]) AS storage) AS storages
+		LEFT JOIN storage_repositories AS sr USING(virtual_storage, relative_path, storage)
+		WHERE r.virtual_storage = $2
+		GROUP BY storages.storage
+		ORDER BY SUM(r.generation - COALESCE(sr.generation, -1))
+		LIMIT 1`
 
 	var newPrimaryStorage string
 	var fallbackChoice bool
diff --git a/internal/praefect/nodes/sql_elector_test.go b/internal/praefect/nodes/sql_elector_test.go
index a47b1fc44..94e5355a6 100644
--- a/internal/praefect/nodes/sql_elector_test.go
+++ b/internal/praefect/nodes/sql_elector_test.go
@@ -266,7 +266,7 @@ func TestElectNewPrimary(t *testing.T) {
 		desc                   string
 		initialReplQueueInsert string
 		expectedPrimary        string
-		defaultChoice          bool
+		fallbackChoice         bool
 	}{
 		{
 			desc: "gitaly-2 storage has more up to date repositories",
@@ -295,12 +295,49 @@ func TestElectNewPrimary(t *testing.T) {
 				('test-shard-0', '/p/5', 'gitaly-2', 5)
 			`,
 			expectedPrimary: "gitaly-2",
-			defaultChoice:   false,
+			fallbackChoice:  false,
+		},
+		{
+			desc: "gitaly-2 storage has less repositories as some may not been replicated yet",
+			initialReplQueueInsert: `
+			INSERT INTO REPOSITORIES
+				(virtual_storage, relative_path, generation)
+			VALUES
+				('test-shard-0', '/p/1', 5),
+				('test-shard-0', '/p/2', 5);
+
+			INSERT INTO STORAGE_REPOSITORIES
+			VALUES
+				('test-shard-0', '/p/1', 'gitaly-1', 5),
+				('test-shard-0', '/p/2', 'gitaly-1', 4),
+				('test-shard-0', '/p/1', 'gitaly-2', 5)`,
+			expectedPrimary: "gitaly-1",
+			fallbackChoice:  false,
+		},
+		{
+			desc: "gitaly-1 is primary as it has less generations behind in total despite it has less repositories",
+			initialReplQueueInsert: `
+			INSERT INTO REPOSITORIES
+				(virtual_storage, relative_path, generation)
+			VALUES
+				('test-shard-0', '/p/1', 2),
+				('test-shard-0', '/p/2', 2),
+				('test-shard-0', '/p/3', 10);
+
+			INSERT INTO STORAGE_REPOSITORIES
+			VALUES
+				('test-shard-0', '/p/2', 'gitaly-1', 1),
+				('test-shard-0', '/p/3', 'gitaly-1', 9),
+				('test-shard-0', '/p/1', 'gitaly-2', 1),
+				('test-shard-0', '/p/2', 'gitaly-2', 1),
+				('test-shard-0', '/p/3', 'gitaly-2', 1)`,
+			expectedPrimary: "gitaly-1",
+			fallbackChoice:  false,
 		},
 		{
 			desc:            "no information about generations results to first candidate",
 			expectedPrimary: "gitaly-1",
-			defaultChoice:   true,
+			fallbackChoice:  true,
 		},
 	}
 
@@ -321,8 +358,8 @@ func TestElectNewPrimary(t *testing.T) {
 			require.NoError(t, err)
 			require.Equal(t, testCase.expectedPrimary, primary.GetStorage())
 
-			defaultChoice := hook.LastEntry().Data["default_choice"].(bool)
-			require.Equal(t, testCase.defaultChoice, defaultChoice)
+			fallbackChoice := hook.LastEntry().Data["fallback_choice"].(bool)
+			require.Equal(t, testCase.fallbackChoice, fallbackChoice)
 		})
 	}
 }