Return more information from GetPartiallyReplicatedRepositories

GetPartiallyReplicatedRepositories returns information about repositories
which have outdated replicas on assigned hosts. The generations returned
are used in `praefect dataloss` to determine whether a repositroy is in
read-only mode or not. With lazy failover, there is no read-only mode
anymore as Praefect can immediately failover to another valid primary.
Praefect doesn't serve reads from outdated replicas, so the repository
would effectively be unavailable if there are no up to date and healthy
replicas. To prepare for updating `praefect dataloss` to account for lazy
failovers, let's return the health status and whether the replica can act
as the primary with each of the replicas. We can later use the ValidPrimary
field to determine if the repository is available and the health status to
ease with debugging why a repository may be unavailable. Other than returning
the additional fields, this commit makes no other behavior changes yet.

2 files changed, 59 insertions, 31 deletions

diff --git a/internal/praefect/datastore/repository_store.go b/internal/praefect/datastore/repository_store.go
index 57a108a64..c975e56cc 100644
--- a/internal/praefect/datastore/repository_store.go
+++ b/internal/praefect/datastore/repository_store.go
@@ -526,6 +526,10 @@ type OutdatedRepositoryStorageDetails struct {
 	BehindBy int
 	// Assigned indicates whether the storage is an assigned host of the repository.
 	Assigned bool
+	// Healthy indicates whether the replica is considered healthy by the consensus of Praefect nodes.
+	Healthy bool
+	// ValidPrimary indicates whether the replica is ready to serve as the primary if necessary.
+	ValidPrimary bool
 }
 
 // OutdatedRepository is a repository with one or more outdated assigned storages.
@@ -545,8 +549,10 @@ func (rs *PostgresRepositoryStore) GetPartiallyReplicatedRepositories(ctx contex
 		return nil, fmt.Errorf("unknown virtual storage: %q", virtualStorage)
 	}
 
-	// The query below gets the generations and assignments of every repository
-	// which has one or more outdated assigned nodes. It works as follows:
+	// The query below gets the status of every repository which has one or more outdated assigned nodes.
+	// The status includes how many changes a replica is behind, whether the replica is assigned host or
+	// not, whether the replica is healthy and whether the replica is considered a valid primary. It works
+	// as follows:
 	//
 	// 1. First we get all the storages which contain the repository from `storage_repositories`. We
 	//    list every copy of the repository as the latest generation could exist on an unassigned
@@ -582,7 +588,9 @@ SELECT
 			json_build_object(
 				'Name', storage,
 				'BehindBy', behind_by,
-				'Assigned', assigned
+				'Assigned', assigned,
+				'Healthy', healthy,
+				'ValidPrimary', valid_primary
 			)
 		)
 	)
@@ -592,7 +600,9 @@ FROM (
 		repositories.primary,
 		storage,
 		repository_generations.generation - COALESCE(storage_repositories.generation, -1) AS behind_by,
-		repository_assignments.storage IS NOT NULL AS assigned
+		repository_assignments.storage IS NOT NULL AS assigned,
+		healthy_storages.storage IS NOT NULL AS healthy,
+		valid_primaries.storage IS NOT NULL AS valid_primary
 	FROM storage_repositories
 	FULL JOIN (
 		SELECT virtual_storage, relative_path, storage
@@ -608,6 +618,8 @@ FROM (
 	) AS repository_assignments USING (virtual_storage, relative_path, storage)
 	JOIN repositories USING (virtual_storage, relative_path)
 	JOIN repository_generations USING (virtual_storage, relative_path)
+	LEFT JOIN healthy_storages USING (virtual_storage, storage)
+	LEFT JOIN valid_primaries USING (virtual_storage, relative_path, storage)
 	WHERE virtual_storage = $1
 	ORDER BY relative_path, "primary", storage
 ) AS outdated_repositories
diff --git a/internal/praefect/datastore/repository_store_test.go b/internal/praefect/datastore/repository_store_test.go
index 7d6a9855d..227acaf93 100644
--- a/internal/praefect/datastore/repository_store_test.go
+++ b/internal/praefect/datastore/repository_store_test.go
@@ -803,6 +803,7 @@ func TestPostgresRepositoryStore_GetPartiallyReplicatedRepositories(t *testing.T
 	for _, tc := range []struct {
 		desc                  string
 		nonExistentRepository bool
+		unhealthyStorages     map[string]struct{}
 		existingGenerations   map[string]int
 		existingAssignments   []string
 		storageDetails        []OutdatedRepositoryStorageDetails
@@ -819,8 +820,8 @@ func TestPostgresRepositoryStore_GetPartiallyReplicatedRepositories(t *testing.T
 			desc:                "unconfigured node contains the latest",
 			existingGenerations: map[string]int{"primary": 0, "secondary-1": 0, "unconfigured": 1},
 			storageDetails: []OutdatedRepositoryStorageDetails{
-				{Name: "primary", BehindBy: 1, Assigned: true},
-				{Name: "secondary-1", BehindBy: 1, Assigned: true},
+				{Name: "primary", BehindBy: 1, Assigned: true, Healthy: true},
+				{Name: "secondary-1", BehindBy: 1, Assigned: true, Healthy: true},
 				{Name: "unconfigured", BehindBy: 0, Assigned: false},
 			},
 		},
@@ -828,32 +829,32 @@ func TestPostgresRepositoryStore_GetPartiallyReplicatedRepositories(t *testing.T
 			desc:                "node has no repository without assignments",
 			existingGenerations: map[string]int{"primary": 0},
 			storageDetails: []OutdatedRepositoryStorageDetails{
-				{Name: "primary", BehindBy: 0, Assigned: true},
-				{Name: "secondary-1", BehindBy: 1, Assigned: true},
+				{Name: "primary", BehindBy: 0, Assigned: true, Healthy: true, ValidPrimary: true},
+				{Name: "secondary-1", BehindBy: 1, Assigned: true, Healthy: true},
 			},
 		},
 		{
 			desc:                "node has outdated repository without assignments",
 			existingGenerations: map[string]int{"primary": 1, "secondary-1": 0},
 			storageDetails: []OutdatedRepositoryStorageDetails{
-				{Name: "primary", BehindBy: 0, Assigned: true},
-				{Name: "secondary-1", BehindBy: 1, Assigned: true},
+				{Name: "primary", BehindBy: 0, Assigned: true, Healthy: true, ValidPrimary: true},
+				{Name: "secondary-1", BehindBy: 1, Assigned: true, Healthy: true},
 			},
 		},
 		{
 			desc:                "node with no repository heavily outdated",
 			existingGenerations: map[string]int{"primary": 10},
 			storageDetails: []OutdatedRepositoryStorageDetails{
-				{Name: "primary", BehindBy: 0, Assigned: true},
-				{Name: "secondary-1", BehindBy: 11, Assigned: true},
+				{Name: "primary", BehindBy: 0, Assigned: true, Healthy: true, ValidPrimary: true},
+				{Name: "secondary-1", BehindBy: 11, Assigned: true, Healthy: true},
 			},
 		},
 		{
 			desc:                "node with a heavily outdated repository",
 			existingGenerations: map[string]int{"primary": 10, "secondary-1": 0},
 			storageDetails: []OutdatedRepositoryStorageDetails{
-				{Name: "primary", BehindBy: 0, Assigned: true},
-				{Name: "secondary-1", BehindBy: 10, Assigned: true},
+				{Name: "primary", BehindBy: 0, Assigned: true, Healthy: true, ValidPrimary: true},
+				{Name: "secondary-1", BehindBy: 10, Assigned: true, Healthy: true},
 			},
 		},
 		{
@@ -876,8 +877,8 @@ func TestPostgresRepositoryStore_GetPartiallyReplicatedRepositories(t *testing.T
 			existingAssignments: []string{"primary", "secondary-1"},
 			existingGenerations: map[string]int{"primary": 0},
 			storageDetails: []OutdatedRepositoryStorageDetails{
-				{Name: "primary", BehindBy: 0, Assigned: true},
-				{Name: "secondary-1", BehindBy: 1, Assigned: true},
+				{Name: "primary", BehindBy: 0, Assigned: true, Healthy: true, ValidPrimary: true},
+				{Name: "secondary-1", BehindBy: 1, Assigned: true, Healthy: true},
 			},
 		},
 		{
@@ -885,8 +886,8 @@ func TestPostgresRepositoryStore_GetPartiallyReplicatedRepositories(t *testing.T
 			existingAssignments: []string{"primary", "secondary-1"},
 			existingGenerations: map[string]int{"primary": 1, "secondary-1": 0},
 			storageDetails: []OutdatedRepositoryStorageDetails{
-				{Name: "primary", BehindBy: 0, Assigned: true},
-				{Name: "secondary-1", BehindBy: 1, Assigned: true},
+				{Name: "primary", BehindBy: 0, Assigned: true, Healthy: true, ValidPrimary: true},
+				{Name: "secondary-1", BehindBy: 1, Assigned: true, Healthy: true},
 			},
 		},
 		{
@@ -894,8 +895,8 @@ func TestPostgresRepositoryStore_GetPartiallyReplicatedRepositories(t *testing.T
 			existingAssignments: []string{"primary"},
 			existingGenerations: map[string]int{"primary": 0, "secondary-1": 1},
 			storageDetails: []OutdatedRepositoryStorageDetails{
-				{Name: "primary", BehindBy: 1, Assigned: true},
-				{Name: "secondary-1", BehindBy: 0, Assigned: false},
+				{Name: "primary", BehindBy: 1, Assigned: true, Healthy: true},
+				{Name: "secondary-1", BehindBy: 0, Assigned: false, Healthy: true, ValidPrimary: true},
 			},
 		},
 		{
@@ -903,8 +904,8 @@ func TestPostgresRepositoryStore_GetPartiallyReplicatedRepositories(t *testing.T
 			existingAssignments: []string{"primary"},
 			existingGenerations: map[string]int{"secondary-1": 0},
 			storageDetails: []OutdatedRepositoryStorageDetails{
-				{Name: "primary", BehindBy: 1, Assigned: true},
-				{Name: "secondary-1", BehindBy: 0, Assigned: false},
+				{Name: "primary", BehindBy: 1, Assigned: true, Healthy: true},
+				{Name: "secondary-1", BehindBy: 0, Assigned: false, Healthy: true, ValidPrimary: true},
 			},
 		},
 		{
@@ -912,7 +913,7 @@ func TestPostgresRepositoryStore_GetPartiallyReplicatedRepositories(t *testing.T
 			existingAssignments: []string{"primary"},
 			existingGenerations: map[string]int{"unconfigured": 0},
 			storageDetails: []OutdatedRepositoryStorageDetails{
-				{Name: "primary", BehindBy: 1, Assigned: true},
+				{Name: "primary", BehindBy: 1, Assigned: true, Healthy: true},
 				{Name: "unconfigured", BehindBy: 0, Assigned: false},
 			},
 		},
@@ -931,8 +932,8 @@ func TestPostgresRepositoryStore_GetPartiallyReplicatedRepositories(t *testing.T
 			existingAssignments: []string{"unconfigured"},
 			existingGenerations: map[string]int{"unconfigured": 0},
 			storageDetails: []OutdatedRepositoryStorageDetails{
-				{Name: "primary", BehindBy: 1, Assigned: true},
-				{Name: "secondary-1", BehindBy: 1, Assigned: true},
+				{Name: "primary", BehindBy: 1, Assigned: true, Healthy: true},
+				{Name: "secondary-1", BehindBy: 1, Assigned: true, Healthy: true},
 				{Name: "unconfigured", BehindBy: 0, Assigned: false},
 			},
 		},
@@ -941,12 +942,27 @@ func TestPostgresRepositoryStore_GetPartiallyReplicatedRepositories(t *testing.T
 			ctx, cancel := testhelper.Context()
 			defer cancel()
 
-			db := getDB(t)
+			tx, err := getDB(t).Begin()
+			require.NoError(t, err)
+			defer tx.Rollback()
 
 			configuredStorages := map[string][]string{"virtual-storage": {"primary", "secondary-1"}}
 
+			var healthyStorages []string
+			for _, storage := range configuredStorages["virtual-storage"] {
+				if _, ok := tc.unhealthyStorages[storage]; ok {
+					continue
+				}
+
+				healthyStorages = append(healthyStorages, storage)
+			}
+
+			testhelper.SetHealthyNodes(t, ctx, tx, map[string]map[string][]string{
+				"praefect-0": {"virtual-storage": healthyStorages},
+			})
+
 			if !tc.nonExistentRepository {
-				_, err := db.ExecContext(ctx, `
+				_, err := tx.ExecContext(ctx, `
 							INSERT INTO repositories (virtual_storage, relative_path, "primary")
 							VALUES ('virtual-storage', 'relative-path', 'repository-primary')
 						`)
@@ -954,26 +970,26 @@ func TestPostgresRepositoryStore_GetPartiallyReplicatedRepositories(t *testing.T
 			}
 
 			for storage, generation := range tc.existingGenerations {
-				_, err := db.ExecContext(ctx, `
+				_, err := tx.ExecContext(ctx, `
 							INSERT INTO storage_repositories VALUES ('virtual-storage', 'relative-path', $1, $2)
 						`, storage, generation)
 				require.NoError(t, err)
 			}
 
 			for _, storage := range tc.existingAssignments {
-				_, err := db.ExecContext(ctx, `
+				_, err := tx.ExecContext(ctx, `
 							INSERT INTO repository_assignments VALUES ('virtual-storage', 'relative-path', $1)
 						`, storage)
 				require.NoError(t, err)
 			}
 
-			_, err := db.ExecContext(ctx, `
+			_, err = tx.ExecContext(ctx, `
 						INSERT INTO shard_primaries (shard_name, node_name, elected_by_praefect, elected_at)
 						VALUES ('virtual-storage', 'virtual-storage-primary', 'ignored', now())
 					`)
 			require.NoError(t, err)
 
-			store := NewPostgresRepositoryStore(db, configuredStorages)
+			store := NewPostgresRepositoryStore(tx, configuredStorages)
 			outdated, err := store.GetPartiallyReplicatedRepositories(ctx, "virtual-storage")
 			require.NoError(t, err)