upgrades: add checkpointing for raftAppliedIndexTermMigration

pkg/upgrade/upgrades/raft_applied_index_term.go

@@ -14,14 +14,18 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
+	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/upgrade"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/retry"
+	"github.com/cockroachdb/errors"
 )
 
 // defaultPageSize controls how many ranges are paged in by default when
@@ -37,12 +41,32 @@ import (
 //   page size of 200 ~ 50s
 const defaultPageSize = 200
 
+// persistWatermarkBatchInterval specifies how often to persist the progress
+// watermark (in batches). 5 batches means we'll checkpoint every 1000 ranges.
+const persistWatermarkBatchInterval = 5
+
 func raftAppliedIndexTermMigration(
-	ctx context.Context, cv clusterversion.ClusterVersion, deps upgrade.SystemDeps, _ *jobs.Job,
+	ctx context.Context, cv clusterversion.ClusterVersion, deps upgrade.SystemDeps, job *jobs.Job,
 ) error {
+	// Fetch the migration's watermark (latest migrated range's end key), in case
+	// we're resuming a previous migration.
+	var resumeWatermark roachpb.RKey
+	if progress, ok := job.Progress().Details.(*jobspb.Progress_Migration); ok {
+		if len(progress.Migration.Watermark) > 0 {
+			resumeWatermark = append(resumeWatermark, progress.Migration.Watermark...)
+			log.Infof(ctx, "loaded migration watermark %s, resuming", resumeWatermark)
+		}
+	}
+
+	retryOpts := retry.Options{
+		InitialBackoff: 100 * time.Millisecond,
+		MaxRetries:     5,
+	}
+
 	var batchIdx, numMigratedRanges int
 	init := func() { batchIdx, numMigratedRanges = 1, 0 }
 	if err := deps.Cluster.IterateRangeDescriptors(ctx, defaultPageSize, init, func(descriptors ...roachpb.RangeDescriptor) error {
+		var progress jobspb.MigrationProgress
 		for _, desc := range descriptors {
 			// NB: This is a bit of a wart. We want to reach the first range,
 			// but we can't address the (local) StartKey. However, keys.LocalMax
@@ -51,8 +75,30 @@ func raftAppliedIndexTermMigration(
 			if bytes.Compare(desc.StartKey, keys.LocalMax) < 0 {
 				start, _ = keys.Addr(keys.LocalMax)
 			}
-			if err := deps.DB.Migrate(ctx, start, end, cv.Version); err != nil {
+
+			// Skip any ranges below the resume watermark.
+			if bytes.Compare(end, resumeWatermark) <= 0 {
+				continue
+			}
+
+			// Migrate the range, with a few retries.
+			if err := retryOpts.Do(ctx, func(ctx context.Context) error {
+				err := deps.DB.Migrate(ctx, start, end, cv.Version)
+				if err != nil {
+					log.Errorf(ctx, "range %d migration failed, retrying: %s", desc.RangeID, err)
+				}
 				return err
+			}); err != nil {
+				return err
+			}
+
+			progress.Watermark = end
+		}
+
+		// Persist the migration's progress.
+		if batchIdx%persistWatermarkBatchInterval == 0 && len(progress.Watermark) > 0 {
+			if err := job.SetProgress(ctx, nil, progress); err != nil {
+				return errors.Wrap(err, "failed to record migration progress")
 			}
 		}