Merge branch 'smh-elector-assignments' into 'master'

Only elect assigned nodes in per repository elector

See merge request gitlab-org/gitaly!2793

2 files changed, 30 insertions, 19 deletions

diff --git a/internal/praefect/nodes/per_repository.go b/internal/praefect/nodes/per_repository.go
index c357eb29c..f30d50073 100644
--- a/internal/praefect/nodes/per_repository.go
+++ b/internal/praefect/nodes/per_repository.go
@@ -92,6 +92,7 @@ UPDATE repositories
 		LEFT JOIN storage_repositories USING (virtual_storage, storage)
 		WHERE virtual_storage = repositories.virtual_storage
 		AND storage_repositories.relative_path = repositories.relative_path
+		AND assigned
 		ORDER BY generation DESC NULLS LAST, random()
 		LIMIT 1
 	)
@@ -100,10 +101,6 @@ WHERE NOT EXISTS (
 	FROM healthy_storages
 	WHERE virtual_storage = repositories.virtual_storage
 	AND storage = repositories."primary"
-) AND EXISTS (
-	SELECT 1
-	FROM healthy_storages
-	WHERE virtual_storage = repositories.virtual_storage
 )`, pq.StringArray(virtualStorages), pq.StringArray(physicalStorages)); err != nil {
 		return fmt.Errorf("query: %w", err)
 	}
diff --git a/internal/praefect/nodes/per_repository_test.go b/internal/praefect/nodes/per_repository_test.go
index 565c05127..d0e82592b 100644
--- a/internal/praefect/nodes/per_repository_test.go
+++ b/internal/praefect/nodes/per_repository_test.go
@@ -19,7 +19,12 @@ func TestPerRepositoryElector(t *testing.T) {
 	ctx, cancel := testhelper.Context()
 	defer cancel()
 
-	type state map[string]map[string]map[string]int
+	type storageRecord struct {
+		generation int
+		assigned   bool
+	}
+
+	type state map[string]map[string]map[string]storageRecord
 
 	type matcher func(t testing.TB, primary string)
 	any := func(expected ...string) matcher {
@@ -52,8 +57,8 @@ func TestPerRepositoryElector(t *testing.T) {
 			state: state{
 				"virtual-storage-1": {
 					"relative-path-1": {
-						"gitaly-1": 1,
-						"gitaly-2": 0,
+						"gitaly-1": {generation: 1, assigned: true},
+						"gitaly-2": {generation: 0, assigned: true},
 					},
 				},
 			},
@@ -71,8 +76,8 @@ func TestPerRepositoryElector(t *testing.T) {
 			state: state{
 				"virtual-storage-1": {
 					"relative-path-1": {
-						"gitaly-1": 1,
-						"gitaly-2": 0,
+						"gitaly-1": {generation: 1, assigned: true},
+						"gitaly-2": {generation: 0, assigned: true},
 					},
 				},
 			},
@@ -102,8 +107,8 @@ func TestPerRepositoryElector(t *testing.T) {
 			state: state{
 				"virtual-storage-1": {
 					"relative-path-1": {
-						"gitaly-1": 0,
-						"gitaly-2": 0,
+						"gitaly-1": {generation: 0, assigned: true},
+						"gitaly-2": {generation: 0, assigned: true},
 					},
 				},
 			},
@@ -121,9 +126,9 @@ func TestPerRepositoryElector(t *testing.T) {
 			state: state{
 				"virtual-storage-1": {
 					"relative-path-1": {
-						"gitaly-1": 1,
-						"gitaly-2": 1,
-						"gitaly-3": 0,
+						"gitaly-1": {generation: 1, assigned: true},
+						"gitaly-2": {generation: 1, assigned: true},
+						"gitaly-3": {generation: 0, assigned: true},
 					},
 				},
 			},
@@ -147,8 +152,8 @@ func TestPerRepositoryElector(t *testing.T) {
 			state: state{
 				"virtual-storage-1": {
 					"relative-path-1": {
-						"gitaly-1": 1,
-						"gitaly-3": 0,
+						"gitaly-1": {generation: 1, assigned: true},
+						"gitaly-3": {generation: 0, assigned: true},
 					},
 				},
 			},
@@ -172,7 +177,8 @@ func TestPerRepositoryElector(t *testing.T) {
 			state: state{
 				"virtual-storage-1": {
 					"relative-path-1": {
-						"gitaly-1": 1,
+						"gitaly-1": {generation: 1, assigned: true},
+						"gitaly-2": {generation: 1, assigned: false},
 					},
 				},
 			},
@@ -204,8 +210,16 @@ func TestPerRepositoryElector(t *testing.T) {
 			rs := datastore.NewPostgresRepositoryStore(db, nil)
 			for virtualStorage, relativePaths := range tc.state {
 				for relativePath, storages := range relativePaths {
-					for storage, generation := range storages {
-						require.NoError(t, rs.SetGeneration(ctx, virtualStorage, relativePath, storage, generation))
+					for storage, record := range storages {
+						require.NoError(t, rs.SetGeneration(ctx, virtualStorage, relativePath, storage, record.generation))
+
+						_, err := db.ExecContext(ctx, `
+							UPDATE storage_repositories SET assigned = $4
+							WHERE virtual_storage = $1
+							AND   relative_path = $2
+							AND   storage = $3
+						`, virtualStorage, relativePath, storage, record.assigned)
+						require.NoError(t, err)
 					}
 				}
 			}