Implement 'praefect verify' subcommandsmh-verify-subcmd

Praefect periodically verifies the repository metadata in the background.
The interval may be too long especially if there was an incident, say
a disk failure. After recovering the Gitaly node, the disk may be completely
empty or may contains an older snapshot which does not contain all expected
repositories. In such cases, it would be great if Praefect could be manually
instructed to verify the storage again as soon as possible rather than waiting
for the next scheduled verification interval to pass. This commit adds the
'praefect verify' subcommand that allows for doing that. It takes in either a
repository id, a virtual storage or a (virtual storage, storage) tuple and
marks all replicas matching the selector as being unverified. Praefect's
metadata verifier will then prioritize verifying these repositories over other
repositories pending verification. This allows administrators to speed up
the verification process and thus recovery.

Changelog: added

6 files changed, 415 insertions, 0 deletions

diff --git a/cmd/praefect/subcmd.go b/cmd/praefect/subcmd.go
index 3a4148244..e62ead65c 100644
--- a/cmd/praefect/subcmd.go
+++ b/cmd/praefect/subcmd.go
@@ -52,6 +52,7 @@ var subcommands = map[string]subcmd{
 		praefect.NewClockSyncCheck(helper.CheckClockSync),
 	),
 	metadataCmdName: newMetadataSubcommand(os.Stdout),
+	verifyCmdName:   newVerifySubcommand(os.Stdout),
 }
 
 // subCommand returns an exit code, to be fed into os.Exit.
diff --git a/cmd/praefect/subcmd_verify.go b/cmd/praefect/subcmd_verify.go
new file mode 100644
index 000000000..340731409
--- /dev/null
+++ b/cmd/praefect/subcmd_verify.go
@@ -0,0 +1,85 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"flag"
+	"fmt"
+	"io"
+
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
+	"gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
+)
+
+const verifyCmdName = "verify"
+
+type verifySubcommand struct {
+	stdout         io.Writer
+	repositoryID   int64
+	virtualStorage string
+	storage        string
+}
+
+func newVerifySubcommand(stdout io.Writer) *verifySubcommand {
+	return &verifySubcommand{stdout: stdout}
+}
+
+func (cmd *verifySubcommand) FlagSet() *flag.FlagSet {
+	fs := flag.NewFlagSet(metadataCmdName, flag.ContinueOnError)
+	fs.Int64Var(&cmd.repositoryID, "repository-id", 0, "the id of the repository to verify")
+	fs.StringVar(&cmd.virtualStorage, "virtual-storage", "", "the virtual storage to verify")
+	fs.StringVar(&cmd.storage, "storage", "", "the storage to verify")
+	return fs
+}
+
+func (cmd *verifySubcommand) Exec(flags *flag.FlagSet, cfg config.Config) error {
+	if flags.NArg() > 0 {
+		return unexpectedPositionalArgsError{Command: flags.Name()}
+	}
+
+	var request gitalypb.MarkUnverifiedRequest
+	switch {
+	case cmd.repositoryID != 0:
+		if cmd.virtualStorage != "" || cmd.storage != "" {
+			return errors.New("virtual storage and storage can't be provided with a repository ID")
+		}
+
+		request.Selector = &gitalypb.MarkUnverifiedRequest_RepositoryId{RepositoryId: cmd.repositoryID}
+	case cmd.storage != "":
+		if cmd.virtualStorage == "" {
+			return errors.New("virtual storage must be passed with storage")
+		}
+
+		request.Selector = &gitalypb.MarkUnverifiedRequest_Storage_{
+			Storage: &gitalypb.MarkUnverifiedRequest_Storage{
+				VirtualStorage: cmd.virtualStorage,
+				Storage:        cmd.storage,
+			},
+		}
+	case cmd.virtualStorage != "":
+		request.Selector = &gitalypb.MarkUnverifiedRequest_VirtualStorage{VirtualStorage: cmd.virtualStorage}
+	default:
+		return errors.New("(repository id), (virtual storage) or (virtual storage, storage) required")
+	}
+
+	nodeAddr, err := getNodeAddress(cfg)
+	if err != nil {
+		return fmt.Errorf("get node address: %w", err)
+	}
+
+	ctx := context.TODO()
+	conn, err := subCmdDial(ctx, nodeAddr, cfg.Auth.Token, defaultDialTimeout)
+	if err != nil {
+		return fmt.Errorf("dial: %w", err)
+	}
+	defer conn.Close()
+
+	response, err := gitalypb.NewPraefectInfoServiceClient(conn).MarkUnverified(ctx, &request)
+	if err != nil {
+		return fmt.Errorf("verify replicas: %w", err)
+	}
+
+	fmt.Fprintf(cmd.stdout, "%d replicas marked unverified\n", response.GetReplicasMarked())
+
+	return nil
+}
diff --git a/cmd/praefect/subcmd_verify_test.go b/cmd/praefect/subcmd_verify_test.go
new file mode 100644
index 000000000..107430c7c
--- /dev/null
+++ b/cmd/praefect/subcmd_verify_test.go
@@ -0,0 +1,236 @@
+package main
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/service/info"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper/testdb"
+)
+
+func TestVerifySubcommand(t *testing.T) {
+	t.Parallel()
+
+	// state contains the asserted state as virtual storage -> relative path -> storage -> verified/not verified.
+	type state map[string]map[string]map[string]bool
+
+	startingState := state{
+		"virtual-storage-1": {
+			"relative-path-1": {
+				"unverified": false,
+				"verified-1": true,
+				"verified-2": true,
+			},
+			"relative-path-2": {
+				"unverified": false,
+				"verified-1": true,
+				"verified-2": true,
+			},
+		},
+		"virtual-storage-2": {
+			"relative-path-1": {
+				"unverified": false,
+				"verified-1": true,
+				"verified-2": true,
+			},
+		},
+	}
+
+	for _, tc := range []struct {
+		desc           string
+		args           []string
+		replicasMarked int
+		error          error
+		expectedState  state
+	}{
+		{
+			desc:          "no selector given",
+			error:         errors.New("(repository id), (virtual storage) or (virtual storage, storage) required"),
+			expectedState: startingState,
+		},
+		{
+			desc:          "virtual storage passed with repository id",
+			args:          []string{"-repository-id=1", "-virtual-storage=virtual-storage"},
+			error:         errors.New("virtual storage and storage can't be provided with a repository ID"),
+			expectedState: startingState,
+		},
+		{
+			desc:          "storage passed with repository id",
+			args:          []string{"-repository-id=1", "-storage=storage"},
+			error:         errors.New("virtual storage and storage can't be provided with a repository ID"),
+			expectedState: startingState,
+		},
+		{
+			desc:          "storage passed without virtual storage",
+			args:          []string{"-storage=storage"},
+			error:         errors.New("virtual storage must be passed with storage"),
+			expectedState: startingState,
+		},
+		{
+			desc:          "no repository matched",
+			args:          []string{"-repository-id=1000"},
+			expectedState: startingState,
+		},
+		{
+			desc:           "scheduled by repository id",
+			args:           []string{"-repository-id=1"},
+			replicasMarked: 2,
+			expectedState: state{
+				"virtual-storage-1": {
+					"relative-path-1": {
+						"unverified": false,
+						"verified-1": false,
+						"verified-2": false,
+					},
+					"relative-path-2": {
+						"unverified": false,
+						"verified-1": true,
+						"verified-2": true,
+					},
+				},
+				"virtual-storage-2": {
+					"relative-path-1": {
+						"unverified": false,
+						"verified-1": true,
+						"verified-2": true,
+					},
+				},
+			},
+		},
+		{
+			desc:           "scheduled by virtual storage",
+			args:           []string{"-virtual-storage=virtual-storage-1"},
+			replicasMarked: 4,
+			expectedState: state{
+				"virtual-storage-1": {
+					"relative-path-1": {
+						"unverified": false,
+						"verified-1": false,
+						"verified-2": false,
+					},
+					"relative-path-2": {
+						"unverified": false,
+						"verified-1": false,
+						"verified-2": false,
+					},
+				},
+				"virtual-storage-2": {
+					"relative-path-1": {
+						"unverified": false,
+						"verified-1": true,
+						"verified-2": true,
+					},
+				},
+			},
+		},
+		{
+			desc:           "scheduled by storage",
+			args:           []string{"-virtual-storage=virtual-storage-1", "-storage=verified-1"},
+			replicasMarked: 2,
+			expectedState: state{
+				"virtual-storage-1": {
+					"relative-path-1": {
+						"unverified": false,
+						"verified-1": false,
+						"verified-2": true,
+					},
+					"relative-path-2": {
+						"unverified": false,
+						"verified-1": false,
+						"verified-2": true,
+					},
+				},
+				"virtual-storage-2": {
+					"relative-path-1": {
+						"unverified": false,
+						"verified-1": true,
+						"verified-2": true,
+					},
+				},
+			},
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			db := testdb.New(t)
+
+			rs := datastore.NewPostgresRepositoryStore(db, nil)
+
+			ln, clean := listenAndServe(t, []svcRegistrar{
+				registerPraefectInfoServer(info.NewServer(config.Config{}, rs, nil, nil, nil)),
+			})
+			defer clean()
+
+			ctx := testhelper.Context(t)
+
+			require.NoError(t,
+				rs.CreateRepository(ctx, 1, "virtual-storage-1", "relative-path-1", "replica-path-1", "unverified", []string{"verified-1", "verified-2"}, nil, false, false),
+			)
+			require.NoError(t,
+				rs.CreateRepository(ctx, 2, "virtual-storage-1", "relative-path-2", "replica-path-2", "unverified", []string{"verified-1", "verified-2"}, nil, false, false),
+			)
+			require.NoError(t,
+				rs.CreateRepository(ctx, 3, "virtual-storage-2", "relative-path-1", "replica-path-3", "unverified", []string{"verified-1", "verified-2"}, nil, false, false),
+			)
+
+			_, err := db.ExecContext(ctx, `
+				UPDATE storage_repositories
+				SET verified_at = now()
+				FROM repositories
+				WHERE repositories.repository_id = storage_repositories.repository_id
+				AND   storage != 'unverified'
+			`)
+			require.NoError(t, err)
+
+			stdout := &bytes.Buffer{}
+			cmd := newVerifySubcommand(stdout)
+
+			fs := cmd.FlagSet()
+			require.NoError(t, fs.Parse(tc.args))
+			err = cmd.Exec(fs, config.Config{SocketPath: ln.Addr().String()})
+			testhelper.RequireGrpcError(t, tc.error, err)
+			if tc.error != nil {
+				return
+			}
+
+			require.Equal(t, fmt.Sprintf("%d replicas marked unverified\n", tc.replicasMarked), stdout.String())
+
+			actualState := state{}
+			rows, err := db.QueryContext(ctx, `
+				SELECT
+					repositories.virtual_storage,
+					repositories.relative_path,
+					storage,
+					verified_at IS NOT NULL as verified
+				FROM repositories
+				JOIN storage_repositories USING (repository_id)
+			`)
+			require.NoError(t, err)
+			defer rows.Close()
+
+			for rows.Next() {
+				var virtualStorage, relativePath, storage string
+				var verified bool
+				require.NoError(t, rows.Scan(&virtualStorage, &relativePath, &storage, &verified))
+
+				if actualState[virtualStorage] == nil {
+					actualState[virtualStorage] = map[string]map[string]bool{}
+				}
+
+				if actualState[virtualStorage][relativePath] == nil {
+					actualState[virtualStorage][relativePath] = map[string]bool{}
+				}
+
+				actualState[virtualStorage][relativePath][storage] = verified
+			}
+
+			require.NoError(t, rows.Err())
+			require.Equal(t, tc.expectedState, actualState)
+		})
+	}
+}
diff --git a/internal/praefect/datastore/repository_store.go b/internal/praefect/datastore/repository_store.go
index cf3dda65e..dc28f228c 100644
--- a/internal/praefect/datastore/repository_store.go
+++ b/internal/praefect/datastore/repository_store.go
@@ -139,6 +139,12 @@ type RepositoryStore interface {
 	GetRepositoryMetadata(ctx context.Context, repositoryID int64) (RepositoryMetadata, error)
 	// GetRepositoryMetadataByPath retrieves a repository's metadata by its virtual path.
 	GetRepositoryMetadataByPath(ctx context.Context, virtualStorage, relativePath string) (RepositoryMetadata, error)
+	// MarkUnverified marks replicas of the repository unverified.
+	MarkUnverified(ctx context.Context, repositoryID int64) (int64, error)
+	// MarkVirtualStorageUnverified marks all replicas on the virtual storage as unverified.
+	MarkVirtualStorageUnverified(ctx context.Context, virtualStorage string) (int64, error)
+	// MarkStorageUnverified marsk all replicas on the storage as unverified.
+	MarkStorageUnverified(ctx context.Context, virtualStorage, storage string) (int64, error)
 }
 
 // PostgresRepositoryStore is a Postgres implementation of RepositoryStore.
@@ -153,6 +159,56 @@ func NewPostgresRepositoryStore(db glsql.Querier, configuredStorages map[string]
 	return &PostgresRepositoryStore{db: db, storages: storages(configuredStorages)}
 }
 
+// MarkUnverified marks replicas of the repository unverified.
+func (rs *PostgresRepositoryStore) MarkUnverified(ctx context.Context, repositoryID int64) (int64, error) {
+	result, err := rs.db.ExecContext(ctx, `
+		UPDATE storage_repositories
+		SET verified_at = NULL
+		WHERE repository_id = $1
+		AND verified_at IS NOT NULL
+	`, repositoryID)
+	if err != nil {
+		return 0, fmt.Errorf("query: %w", err)
+	}
+
+	return result.RowsAffected()
+}
+
+// MarkVirtualStorageUnverified marks all replicas on the virtual storage as unverified.
+func (rs *PostgresRepositoryStore) MarkVirtualStorageUnverified(ctx context.Context, virtualStorage string) (int64, error) {
+	result, err := rs.db.ExecContext(ctx, `
+		UPDATE storage_repositories
+		SET verified_at = NULL
+		FROM repositories
+		WHERE repositories.virtual_storage = $1
+		AND   repositories.repository_id   = storage_repositories.repository_id
+		AND   verified_at IS NOT NULL
+	`, virtualStorage)
+	if err != nil {
+		return 0, fmt.Errorf("query: %w", err)
+	}
+
+	return result.RowsAffected()
+}
+
+// MarkStorageUnverified marsk all replicas on the storage as unverified.
+func (rs *PostgresRepositoryStore) MarkStorageUnverified(ctx context.Context, virtualStorage, storage string) (int64, error) {
+	result, err := rs.db.ExecContext(ctx, `
+		UPDATE storage_repositories
+		SET verified_at = NULL
+		FROM repositories
+		WHERE repositories.repository_id = storage_repositories.repository_id
+		AND   repositories.virtual_storage = $1
+		AND   storage_repositories.storage = $2
+		AND   verified_at IS NOT NULL
+	`, virtualStorage, storage)
+	if err != nil {
+		return 0, fmt.Errorf("query: %w", err)
+	}
+
+	return result.RowsAffected()
+}
+
 //nolint: revive,stylecheck // This is unintentionally missing documentation.
 func (rs *PostgresRepositoryStore) GetGeneration(ctx context.Context, repositoryID int64, storage string) (int, error) {
 	const q = `
diff --git a/internal/praefect/datastore/repository_store_mock.go b/internal/praefect/datastore/repository_store_mock.go
index 0770331bd..2a021661d 100644
--- a/internal/praefect/datastore/repository_store_mock.go
+++ b/internal/praefect/datastore/repository_store_mock.go
@@ -5,6 +5,7 @@ import "context"
 // MockRepositoryStore allows for mocking a RepositoryStore by parametrizing its behavior. All methods
 // default to what could be considered success if not set.
 type MockRepositoryStore struct {
+	RepositoryStore
 	GetGenerationFunc                       func(ctx context.Context, repositoryID int64, storage string) (int, error)
 	IncrementGenerationFunc                 func(ctx context.Context, repositoryID int64, primary string, secondaries []string) error
 	GetReplicatedGenerationFunc             func(ctx context.Context, repositoryID int64, source, target string) (int, error)
diff --git a/internal/praefect/service/info/verification.go b/internal/praefect/service/info/verification.go
new file mode 100644
index 000000000..9ef07a129
--- /dev/null
+++ b/internal/praefect/service/info/verification.go
@@ -0,0 +1,36 @@
+package info
+
+import (
+	"context"
+	"fmt"
+
+	"gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
+)
+
+// MarkUnverified marks replicas as unverified. See the protobuf declarations for details.
+func (s *Server) MarkUnverified(ctx context.Context, req *gitalypb.MarkUnverifiedRequest) (*gitalypb.MarkUnverifiedResponse, error) {
+	var markUnverified func() (int64, error)
+	switch selector := req.GetSelector().(type) {
+	case *gitalypb.MarkUnverifiedRequest_RepositoryId:
+		markUnverified = func() (int64, error) {
+			return s.rs.MarkUnverified(ctx, selector.RepositoryId)
+		}
+	case *gitalypb.MarkUnverifiedRequest_VirtualStorage:
+		markUnverified = func() (int64, error) {
+			return s.rs.MarkVirtualStorageUnverified(ctx, selector.VirtualStorage)
+		}
+	case *gitalypb.MarkUnverifiedRequest_Storage_:
+		markUnverified = func() (int64, error) {
+			return s.rs.MarkStorageUnverified(ctx, selector.Storage.VirtualStorage, selector.Storage.Storage)
+		}
+	default:
+		return nil, fmt.Errorf("unknown selector: %T", selector)
+	}
+
+	replicasMarked, err := markUnverified()
+	if err != nil {
+		return nil, fmt.Errorf("schedule verification: %w", err)
+	}
+
+	return &gitalypb.MarkUnverifiedResponse{ReplicasMarked: replicasMarked}, nil
+}