restore: prevent scatter of non-empty ranges in online restore

In conventional restore, we prevent the scattering of non-empty ranges
by passing in a max size for the `AdminScatter` request. In online
restore, while we expect to almost always be scattering empty ranges,
due to the fact that the link phase can retry, it is possible to scatter
ranges that contain external SSTs. This breaks our golden rule of only
scattering empty ranges. This commit teaches online restore to only ever
scatter empty ranges.

Fixes: #144599

Release note: None

Author: kev-cao

pkg/backup/restore_online.go

@@ -9,6 +9,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"strings"
 	"sync/atomic"
 	"time"
 
@@ -304,6 +305,7 @@ func sendAddRemoteSSTWorker(
 	approxDataSize *int64,
 ) func(context.Context) error {
 	return func(ctx context.Context) error {
+		testingKnobs := execCtx.ExecCfg().BackupRestoreTestingKnobs
 		for entry := range restoreSpanEntriesCh {
 			log.Dev.VInfof(ctx, 1, "starting restore of backed up span %s containing %d files", entry.Span, len(entry.Files))
 
@@ -343,6 +345,12 @@ func sendAddRemoteSSTWorker(
 				if err := sendRemoteAddSSTable(ctx, execCtx, file, entry.ElidedPrefix, fromSystemTenant); err != nil {
 					return err
 				}
+
+				if testingKnobs != nil && testingKnobs.AfterAddRemoteSST != nil {
+					if err := testingKnobs.AfterAddRemoteSST(); err != nil {
+						return err
+					}
+				}
 			}
 
 			var rows, dataSize int64
@@ -388,8 +396,16 @@ func sendAdminScatter(
 	ctx, sp := tracing.ChildSpan(ctx, "backup.sendAdminScatter")
 	defer sp.Finish()
 
-	const maxSize = 4 << 20
+	// To obey the golden rule of never scattering non-empty ranges, we set a
+	// maxSize of 1 here. The link phase should generally only ever be scattering
+	// empty ranges, but since the link phase can retry, we may end up attempting
+	// to scatter a range that contains external SSTs. If a scatter request fails
+	// due to the range being non-empty, we can ignore the error.
+	const maxSize = 1
 	_, err := execCtx.ExecCfg().DB.AdminScatter(ctx, scatterKey, maxSize)
+	if err == nil || strings.Contains(err.Error(), "existing range size") {
+		return nil
+	}
 	return err
 }
 

pkg/backup/restore_online_test.go

@@ -21,6 +21,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cloud/nodelocal"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security/securitytest"
 	"github.com/cockroachdb/cockroach/pkg/sql"
@@ -698,6 +700,98 @@ func TestOnlineRestoreDownloadRetryReset(t *testing.T) {
 	require.Equal(t, maxDownloadAttempts*2, attemptCount)
 }
 
+func TestOnlineRestoreFailScatterNonEmptyRanges(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	defer nodelocal.ReplaceNodeLocalForTesting(t.TempDir())()
+
+	// This test first runs online restore and waits for an external SST to be
+	// added to ensure at least one range has ingested an external SST. It then
+	// pauses the job and resumes it. During this second phase, we count the
+	// number of KV scatter requets that are sent and the number of successful
+	// scatter requests. Since at least one range contains an external SST, the
+	// number of scatter requests should be greater than the number of successful
+	// scatter requests.
+	// NB: This relies on the fact that the `TestingResponseFilter` hook is not
+	// called if an admin scatter fails due to the size limit being hit.
+	reachedPause := make(chan struct{})
+	defer close(reachedPause)
+	pauseCh := make(chan struct{})
+	defer close(pauseCh)
+
+	var resumed atomic.Bool
+	var postPauseScatterRequests atomic.Int32
+	var postPauseSuccessfulScatters atomic.Int32
+
+	params := orParams
+	params.ServerArgs.Knobs = base.TestingKnobs{
+		BackupRestore: &sql.BackupRestoreTestingKnobs{
+			AfterAddRemoteSST: func() error {
+				if resumed.Load() {
+					return nil
+				}
+				reachedPause <- struct{}{}
+				<-pauseCh
+				return nil
+			},
+		},
+		Store: &kvserver.StoreTestingKnobs{
+			TestingRequestFilter: func(ctx context.Context, ba *kvpb.BatchRequest) *kvpb.Error {
+				if !resumed.Load() {
+					return nil
+				}
+				for _, req := range ba.Requests {
+					if req.GetInner().Method() == kvpb.AdminScatter {
+						postPauseScatterRequests.Add(1)
+					}
+				}
+				return nil
+			},
+			TestingResponseFilter: func(
+				ctx context.Context, req *kvpb.BatchRequest, resp *kvpb.BatchResponse,
+			) *kvpb.Error {
+				if !resumed.Load() {
+					return nil
+				}
+				for _, r := range req.Requests {
+					if r.GetInner().Method() == kvpb.AdminScatter {
+						postPauseSuccessfulScatters.Add(1)
+					}
+				}
+				return nil
+			},
+		},
+	}
+
+	const numAccounts = 100
+	_, sqlDB, _, cleanupFn := backupRestoreTestSetupWithParams(
+		t, singleNode, numAccounts, InitManualReplication, params,
+	)
+	defer cleanupFn()
+
+	externalStorage := "nodelocal://1/backup"
+	sqlDB.Exec(t, fmt.Sprintf("BACKUP INTO '%s'", externalStorage))
+
+	var linkJobID jobspb.JobID
+	sqlDB.QueryRow(
+		t,
+		`RESTORE DATABASE data FROM LATEST IN $1
+		WITH EXPERIMENTAL DEFERRED COPY, DETACHED, new_db_name='data2'`,
+		externalStorage,
+	).Scan(&linkJobID)
+
+	<-reachedPause
+	sqlDB.Exec(t, "PAUSE JOB $1", linkJobID)
+	jobutils.WaitForJobToPause(t, sqlDB, linkJobID)
+	pauseCh <- struct{}{}
+
+	resumed.Store(true)
+	sqlDB.Exec(t, "RESUME JOB $1", linkJobID)
+	jobutils.WaitForJobToSucceed(t, sqlDB, linkJobID)
+
+	require.Greater(t, postPauseScatterRequests.Load(), postPauseSuccessfulScatters.Load())
+}
+
 func bankOnlineRestore(
 	t *testing.T,
 	sqlDB *sqlutils.SQLRunner,

pkg/sql/exec_util_backup.go

@@ -86,6 +86,10 @@ type BackupRestoreTestingKnobs struct {
 	// RunBeforeDownloadCleanup is called before we cleanup after all external
 	// files have been download.
 	RunBeforeDownloadCleanup func() error
+
+	// AfterAddRemoteSST is called after a remote SST is linked to pebble during
+	// the link phase of online restore.
+	AfterAddRemoteSST func() error
 }
 
 var _ base.ModuleTestingKnobs = &BackupRestoreTestingKnobs{}