read-only metrics for repository specific primaries

Adds support for collecting read-only repository metrics taking
per repository primaries in to account. The new functionality is
behind a boolean flag in order to keep metric collection working for
virtual storage primaries as well. In-memory implementations have been
removed as there is no support for per repository primaries in the
local elector.

5 files changed, 190 insertions, 215 deletions

diff --git a/internal/praefect/datastore/collector.go b/internal/praefect/datastore/collector.go
index e093535bc..1dc7654a3 100644
--- a/internal/praefect/datastore/collector.go
+++ b/internal/praefect/datastore/collector.go
@@ -2,9 +2,12 @@ package datastore
 
 import (
 	"context"
+	"fmt"
 
+	"github.com/lib/pq"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/sirupsen/logrus"
+	"gitlab.com/gitlab-org/gitaly/internal/praefect/datastore/glsql"
 )
 
 var descReadOnlyRepositories = prometheus.NewDesc(
@@ -14,25 +17,22 @@ var descReadOnlyRepositories = prometheus.NewDesc(
 	nil,
 )
 
-// PrimaryGetter is an interface used by RepositoryStoreCollector to
-// get information of primary assignments of the configured virtual
-// storages.
-type PrimaryGetter interface {
-	// GetPrimaries returns primaries by their virtual storages. If a virtual
-	// storage does not have a primary, its entry should be an empty string.
-	GetPrimaries(ctx context.Context) (map[string]string, error)
-}
-
 // RepositoryStoreCollector collects metrics from the RepositoryStore.
 type RepositoryStoreCollector struct {
-	log logrus.FieldLogger
-	rs  RepositoryStore
-	pg  PrimaryGetter
+	log              logrus.FieldLogger
+	db               glsql.Querier
+	virtualStorages  []string
+	repositoryScoped bool
 }
 
 // NewRepositoryStoreCollector returns a new collector.
-func NewRepositoryStoreCollector(log logrus.FieldLogger, rs RepositoryStore, pg PrimaryGetter) *RepositoryStoreCollector {
-	return &RepositoryStoreCollector{log, rs, pg}
+func NewRepositoryStoreCollector(log logrus.FieldLogger, virtualStorages []string, db glsql.Querier, repositoryScoped bool) *RepositoryStoreCollector {
+	return &RepositoryStoreCollector{
+		log:              log.WithField("component", "RepositoryStoreCollector"),
+		db:               db,
+		virtualStorages:  virtualStorages,
+		repositoryScoped: repositoryScoped,
+	}
 }
 
 func (c *RepositoryStoreCollector) Describe(ch chan<- *prometheus.Desc) {
@@ -40,19 +40,86 @@ func (c *RepositoryStoreCollector) Describe(ch chan<- *prometheus.Desc) {
 }
 
 func (c *RepositoryStoreCollector) Collect(ch chan<- prometheus.Metric) {
-	vsPrimaries, err := c.pg.GetPrimaries(context.TODO())
+	readOnlyCounts, err := c.queryMetrics(context.TODO())
 	if err != nil {
-		c.log.WithError(err).Error("RepositoryStoreCollector: get virtual storage primaries")
+		c.log.WithError(err).Error("failed collecting read-only repository count metric")
 		return
 	}
 
-	readOnlyCounts, err := c.rs.CountReadOnlyRepositories(context.TODO(), vsPrimaries)
+	for _, vs := range c.virtualStorages {
+		ch <- prometheus.MustNewConstMetric(descReadOnlyRepositories, prometheus.GaugeValue, float64(readOnlyCounts[vs]), vs)
+	}
+}
+
+// queryMetrics queries the number of read-only repositories from the database.
+// A repository is in read-only mode when its primary storage is not on the latest
+// generation.
+//
+// There are two variants on the query:
+//
+// 1. virtualStorageScopedQuery considers virtual storage scoped primaries. Virtual storage's
+//    primary is stored in the `shard_primaries` table in the `node_name` column.
+// 2. repositoryScopedQuery considers repository specific primaries. Repository scoped
+//    primaries are stored in the `primary` column in the `repositories` table.
+//
+// Both queries cross-reference the `repositories` and `storage_repositories` tables
+// to see if the primary storage of the repository is on the latest generation. If not,
+// it's added to the returned count.
+//
+// The query operating on virtual storage scoped primaries will be dropped once the migration
+// to repository scoped primaries is finished.
+func (c *RepositoryStoreCollector) queryMetrics(ctx context.Context) (map[string]int, error) {
+	const virtualStorageScopedQuery = `
+SELECT repositories.virtual_storage, COUNT(*)
+FROM repositories
+LEFT JOIN shard_primaries ON
+	shard_primaries.shard_name = repositories.virtual_storage AND
+	shard_primaries.demoted = false
+LEFT JOIN storage_repositories ON
+	repositories.virtual_storage = storage_repositories.virtual_storage AND
+	repositories.relative_path = storage_repositories.relative_path AND
+	shard_primaries.node_name = storage_repositories.storage
+WHERE 
+	COALESCE(storage_repositories.generation, -1) < repositories.generation AND
+	repositories.virtual_storage = ANY($1)
+GROUP BY repositories.virtual_storage;
+	`
+
+	const repositoryScopedQuery = `
+SELECT repositories.virtual_storage, COUNT(*)
+FROM repositories
+LEFT JOIN storage_repositories ON 
+	repositories.virtual_storage = storage_repositories.virtual_storage AND
+	repositories.relative_path = storage_repositories.relative_path AND
+	repositories.primary = storage_repositories.storage
+WHERE
+	COALESCE(storage_repositories.generation, -1) < repositories.generation AND
+	repositories.virtual_storage = ANY($1)
+GROUP BY repositories.virtual_storage
+`
+
+	query := virtualStorageScopedQuery
+	if c.repositoryScoped {
+		query = repositoryScopedQuery
+	}
+
+	rows, err := c.db.QueryContext(ctx, query, pq.StringArray(c.virtualStorages))
 	if err != nil {
-		c.log.WithError(err).Error("RepositoryStoreCollector: count read-only repositories")
-		return
+		return nil, fmt.Errorf("query: %w", err)
 	}
+	defer rows.Close()
 
-	for virtualStorage, readOnlyCount := range readOnlyCounts {
-		ch <- prometheus.MustNewConstMetric(descReadOnlyRepositories, prometheus.GaugeValue, float64(readOnlyCount), virtualStorage)
+	vsReadOnly := make(map[string]int)
+	for rows.Next() {
+		var vs string
+		var count int
+
+		if err := rows.Scan(&vs, &count); err != nil {
+			return nil, fmt.Errorf("scan: %w", err)
+		}
+
+		vsReadOnly[vs] = count
 	}
+
+	return vsReadOnly, rows.Err()
 }
diff --git a/internal/praefect/datastore/collector_test.go b/internal/praefect/datastore/collector_test.go
index 0c5f022e5..3c521dcc4 100644
--- a/internal/praefect/datastore/collector_test.go
+++ b/internal/praefect/datastore/collector_test.go
@@ -1,46 +1,123 @@
+// +build postgres
+
 package datastore
 
 import (
-	"context"
+	"fmt"
 	"strings"
 	"testing"
 
 	"github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
 )
 
-type primaryGetterFunc func(context.Context) (map[string]string, error)
-
-func (pg primaryGetterFunc) GetPrimaries(ctx context.Context) (map[string]string, error) {
-	return pg(ctx)
-}
-
 func TestRepositoryStoreCollector(t *testing.T) {
 	ctx, cancel := testhelper.Context()
 	defer cancel()
 
-	rs := NewMemoryRepositoryStore(nil)
-	require.NoError(t, rs.SetGeneration(ctx, "some-read-only", "read-only", "secondary", 1))
-	require.NoError(t, rs.SetGeneration(ctx, "some-read-only", "read-only", "primary", 0))
-	require.NoError(t, rs.SetGeneration(ctx, "no-records", "writable", "primary", 0))
-	require.NoError(t, rs.SetGeneration(ctx, "no-primary", "read-only", "primary", 0))
-
-	c := NewRepositoryStoreCollector(nil, rs, primaryGetterFunc(func(context.Context) (map[string]string, error) {
-		return map[string]string{
-			"some-read-only": "primary",
-			"all-writable":   "primary",
-			"no-records":     "primary",
-			"no-primary":     "",
-		}, nil
-	}))
-
-	require.NoError(t, testutil.CollectAndCompare(c, strings.NewReader(`
+	db := getDB(t)
+	rs := NewPostgresRepositoryStore(db, nil)
+
+	state := map[string]map[string]map[string]int{
+		"some-read-only": {
+			"read-only": {
+				"vs-primary":   0,
+				"repo-primary": 0,
+				"secondary":    1,
+			},
+			"writable": {
+				"vs-primary":   1,
+				"repo-primary": 1,
+				"secondary":    1,
+			},
+			"repo-writable": {
+				"vs-primary":   0,
+				"repo-primary": 1,
+				"secondary":    1,
+			},
+		},
+		"all-writable": {
+			"writable": {
+				"vs-primary":   0,
+				"repo-primary": 0,
+			},
+		},
+		"unconfigured": {
+			"read-only": {
+				"secondary": 1,
+			},
+		},
+		"no-records": {},
+		"no-primary": {
+			"read-only": {
+				"secondary": 0,
+			},
+		},
+	}
+	for virtualStorage, relativePaths := range state {
+		demoted := false
+		if virtualStorage == "no-primary" {
+			demoted = true
+		}
+		_, err := db.ExecContext(ctx, `
+			INSERT INTO shard_primaries (shard_name, node_name, elected_by_praefect, elected_at, demoted)
+			VALUES ($1, 'vs-primary', 'not-needed', now(), $2)
+			`, virtualStorage, demoted,
+		)
+		require.NoError(t, err)
+
+		for relativePath, storages := range relativePaths {
+			if virtualStorage != "no-primary" {
+				_, err := db.ExecContext(ctx, `
+						INSERT INTO repositories (virtual_storage, relative_path, "primary")
+						VALUES ($1, $2, 'repo-primary')
+						`, virtualStorage, relativePath,
+				)
+				require.NoError(t, err)
+			}
+
+			for storage, generation := range storages {
+				require.NoError(t, rs.SetGeneration(ctx, virtualStorage, relativePath, storage, generation))
+			}
+		}
+	}
+
+	var virtualStorages []string
+	for vs := range state {
+		if vs == "unconfigured" {
+			continue
+		}
+
+		virtualStorages = append(virtualStorages, vs)
+	}
+
+	for _, tc := range []struct {
+		desc              string
+		repositoryScoped  bool
+		someReadOnlyCount int
+	}{
+		{
+			desc:              "repository scoped",
+			someReadOnlyCount: 1,
+			repositoryScoped:  true,
+		},
+		{
+			desc:              "virtual storage scoped",
+			someReadOnlyCount: 2,
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			c := NewRepositoryStoreCollector(logrus.New(), virtualStorages, db, tc.repositoryScoped)
+			require.NoError(t, testutil.CollectAndCompare(c, strings.NewReader(fmt.Sprintf(`
 # HELP gitaly_praefect_read_only_repositories Number of repositories in read-only mode within a virtual storage.
 # TYPE gitaly_praefect_read_only_repositories gauge
-gitaly_praefect_read_only_repositories{virtual_storage="some-read-only"} 1
 gitaly_praefect_read_only_repositories{virtual_storage="all-writable"} 0
 gitaly_praefect_read_only_repositories{virtual_storage="no-records"} 0
 gitaly_praefect_read_only_repositories{virtual_storage="no-primary"} 1
-`)))
+gitaly_praefect_read_only_repositories{virtual_storage="some-read-only"} %d
+`, tc.someReadOnlyCount))))
+		})
+	}
 }
diff --git a/internal/praefect/datastore/repository_memory.go b/internal/praefect/datastore/repository_memory.go
index cb4216f7b..61b2a691e 100644
--- a/internal/praefect/datastore/repository_memory.go
+++ b/internal/praefect/datastore/repository_memory.go
@@ -254,26 +254,6 @@ func (m *MemoryRepositoryStore) GetOutdatedRepositories(ctx context.Context, vir
 	return outdatedRepos, nil
 }
 
-func (m *MemoryRepositoryStore) CountReadOnlyRepositories(ctx context.Context, vsPrimaries map[string]string) (map[string]int, error) {
-	m.m.Lock()
-	defer m.m.Unlock()
-
-	vsReadOnly := make(map[string]int, len(vsPrimaries))
-	for vs, primary := range vsPrimaries {
-		vsReadOnly[vs] = 0
-		relativePaths := m.virtualStorageState[vs]
-		for relativePath := range relativePaths {
-			expectedGeneration := m.getRepositoryGeneration(vs, relativePath)
-			actualGeneration := m.getStorageGeneration(vs, relativePath, primary)
-			if actualGeneration < expectedGeneration {
-				vsReadOnly[vs]++
-			}
-		}
-	}
-
-	return vsReadOnly, nil
-}
-
 func (m *MemoryRepositoryStore) getRepositoryGeneration(virtualStorage, relativePath string) int {
 	generations, ok := m.storageState[virtualStorage][relativePath]
 	if !ok {
diff --git a/internal/praefect/datastore/repository_memory_test.go b/internal/praefect/datastore/repository_memory_test.go
index 161b63d4d..762768f21 100644
--- a/internal/praefect/datastore/repository_memory_test.go
+++ b/internal/praefect/datastore/repository_memory_test.go
@@ -524,98 +524,4 @@ func testRepositoryStore(t *testing.T, newStore repositoryStoreFactory) {
 			})
 		}
 	})
-
-	t.Run("CountReadOnlyRepositories", func(t *testing.T) {
-		rs, requireState := newStore(t, nil)
-
-		t.Run("no read-only repositories", func(t *testing.T) {
-			counts, err := rs.CountReadOnlyRepositories(ctx, map[string]string{
-				"virtual-storage-1": "primary-1",
-				"virtual-storage-2": "primary-2",
-			})
-			require.NoError(t, err)
-			require.Equal(t, map[string]int{
-				"virtual-storage-1": 0,
-				"virtual-storage-2": 0,
-			}, counts)
-		})
-
-		t.Run("read-only repositories", func(t *testing.T) {
-			require.NoError(t, rs.SetGeneration(ctx, "some-read-only", "read-only-outdated", "secondary", 1))
-			require.NoError(t, rs.SetGeneration(ctx, "some-read-only", "read-only-outdated", "primary", 0))
-			require.NoError(t, rs.SetGeneration(ctx, "some-read-only", "read-only-no-record", "secondary", 0))
-			require.NoError(t, rs.SetGeneration(ctx, "some-read-only", "writable", "secondary", 0))
-			require.NoError(t, rs.SetGeneration(ctx, "some-read-only", "writable", "primary", 0))
-			require.NoError(t, rs.SetGeneration(ctx, "all-writable", "writable", "secondary", 0))
-			require.NoError(t, rs.SetGeneration(ctx, "all-writable", "writable", "primary", 0))
-
-			requireState(t, ctx,
-				virtualStorageState{
-					"some-read-only": {
-						"read-only-outdated":  struct{}{},
-						"read-only-no-record": struct{}{},
-						"writable":            struct{}{},
-					},
-					"all-writable": {
-						"writable": struct{}{},
-					},
-				},
-				storageState{
-					"some-read-only": {
-						"read-only-outdated": {
-							"secondary": 1,
-							"primary":   0,
-						},
-						"read-only-no-record": {
-							"secondary": 0,
-						},
-						"writable": {
-							"secondary": 0,
-							"primary":   0,
-						},
-					},
-					"all-writable": {
-						"writable": {
-							"secondary": 0,
-							"primary":   0,
-						},
-					},
-				},
-			)
-
-			t.Run("primaries", func(t *testing.T) {
-				counts, err := rs.CountReadOnlyRepositories(ctx, map[string]string{
-					"some-read-only": "primary",
-					"all-writable":   "primary",
-					"no-records":     "primary",
-				})
-				require.NoError(t, err)
-				require.Equal(t, map[string]int{
-					"some-read-only": 2,
-					"all-writable":   0,
-					"no-records":     0,
-				}, counts)
-			})
-
-			t.Run("no primaries", func(t *testing.T) {
-				counts, err := rs.CountReadOnlyRepositories(ctx, map[string]string{
-					"some-read-only": "",
-					"all-writable":   "",
-					"no-records":     "",
-				})
-				require.NoError(t, err)
-				require.Equal(t, map[string]int{
-					"some-read-only": 3,
-					"all-writable":   1,
-					"no-records":     0,
-				}, counts)
-			})
-
-			t.Run("no virtual storages", func(t *testing.T) {
-				counts, err := rs.CountReadOnlyRepositories(ctx, nil)
-				require.NoError(t, err)
-				require.Equal(t, map[string]int{}, counts)
-			})
-		})
-	})
 }
diff --git a/internal/praefect/datastore/repository_postgres.go b/internal/praefect/datastore/repository_postgres.go
index e00974f3b..7f9805084 100644
--- a/internal/praefect/datastore/repository_postgres.go
+++ b/internal/praefect/datastore/repository_postgres.go
@@ -83,10 +83,6 @@ type RepositoryStore interface {
 	// with key structure `relative_path-> storage -> generation`, indicating how many changes a storage is missing for a given
 	// repository.
 	GetOutdatedRepositories(ctx context.Context, virtualStorage string) (map[string]map[string]int, error)
-	// CountReadOnlyRepositories returns the number of read-only repositories within each virtual storage. Takes in a map
-	// of configured virtual storages with their primary nodes. The primary of a virtual storage may be an empty string
-	// if there is no assigned primary. In such cases, every repository on the virtual storage is counted as read-only.
-	CountReadOnlyRepositories(ctx context.Context, virtualStoragePrimaries map[string]string) (map[string]int, error)
 }
 
 // PostgresRepositoryStore is a Postgres implementation of RepositoryStore.
@@ -466,54 +462,3 @@ WHERE COALESCE(storage_repositories.generation, -1) < expected_repositories.gene
 
 	return outdated, rows.Err()
 }
-
-func (rs *PostgresRepositoryStore) CountReadOnlyRepositories(ctx context.Context, virtualStoragePrimaries map[string]string) (map[string]int, error) {
-	const q = `
-SELECT virtual_storage, COUNT(*)
-FROM (
-	SELECT virtual_storage, relative_path, MAX(storage_repositories.generation) AS generation
-	FROM repositories
-	JOIN storage_repositories USING (virtual_storage, relative_path)
-	GROUP BY virtual_storage, relative_path
-) AS expected_repositories
-JOIN (
-	SELECT
-		unnest($1::text[]) AS virtual_storage,
-		unnest($2::text[]) AS storage
-) AS primaries USING (virtual_storage)
-LEFT JOIN storage_repositories USING (virtual_storage, relative_path, storage)
-WHERE COALESCE(storage_repositories.generation, -1) < expected_repositories.generation
-GROUP BY virtual_storage
-	`
-
-	virtualStorages := make([]string, 0, len(virtualStoragePrimaries))
-	primaries := make([]string, 0, len(virtualStoragePrimaries))
-	for virtualStorage, primary := range virtualStoragePrimaries {
-		virtualStorages = append(virtualStorages, virtualStorage)
-		primaries = append(primaries, primary)
-	}
-
-	rows, err := rs.db.QueryContext(ctx, q, pq.StringArray(virtualStorages), pq.StringArray(primaries))
-	if err != nil {
-		return nil, err
-	}
-	defer rows.Close()
-
-	vsReadOnly := make(map[string]int, len(virtualStorages))
-	for virtualStorage := range virtualStoragePrimaries {
-		vsReadOnly[virtualStorage] = 0
-	}
-
-	for rows.Next() {
-		var vs string
-		var count int
-
-		if err := rows.Scan(&vs, &count); err != nil {
-			return nil, err
-		}
-
-		vsReadOnly[vs] = count
-	}
-
-	return vsReadOnly, rows.Err()
-}