backup: add retry loop to doDownloadFiles

Previously, `doDownloadFiles` would immediately fail the download job
upon encountering an error. This commit teaches the download workers to
retry the download requests before giving up and failing.

Epic: CRDB-51394

Fixes: #148081

Release note: Download phase of fast restore now will retry downloads
before giving up.

Author: kev-cao

pkg/backup/restore_online.go

@@ -64,6 +64,7 @@ var onlineRestoreLayerLimit = settings.RegisterIntSetting(
 )
 
 const linkCompleteKey = "link_complete"
+const maxDownloadAttempts = 5
 
 // splitAndScatter runs through all entries produced by genSpans splitting and
 // scattering the key-space designated by the passed rewriter such that if all
@@ -556,15 +557,34 @@ func (r *restoreResumer) sendDownloadWorker(
 		ctx, tsp := tracing.ChildSpan(ctx, "backup.sendDownloadWorker")
 		defer tsp.Finish()
 
-		for rt := retry.StartWithCtx(
-			ctx, retry.Options{InitialBackoff: time.Millisecond * 100, MaxBackoff: time.Second * 10},
-		); ; rt.Next() {
+		testingKnobs := execCtx.ExecCfg().BackupRestoreTestingKnobs
+		for {
 			if err := ctx.Err(); err != nil {
 				return err
 			}
 
-			if err := sendDownloadSpan(ctx, execCtx, spans); err != nil {
-				return err
+			var err error
+			for r := retry.StartWithCtx(ctx, retry.Options{
+				InitialBackoff: time.Millisecond * 100,
+				MaxBackoff:     time.Second,
+				MaxRetries:     maxDownloadAttempts - 1,
+			}); r.Next(); {
+				err = func() error {
+					if testingKnobs != nil && testingKnobs.RunBeforeSendingDownloadSpan != nil {
+						if err := testingKnobs.RunBeforeSendingDownloadSpan(); err != nil {
+							return err
+						}
+					}
+					return sendDownloadSpan(ctx, execCtx, spans)
+				}()
+				if err == nil {
+					break
+				}
+				log.VInfof(ctx, 1, "attempt %d failed to download spans: %v", r.CurrentAttempt(), err)
+			}
+
+			if err != nil {
+				return errors.Wrapf(err, "retries exhausted for sending download spans")
 			}
 
 			// Wait for the completion poller to signal that it has checked our work.
@@ -576,6 +596,11 @@ func (r *restoreResumer) sendDownloadWorker(
 			case <-ctx.Done():
 				return ctx.Err()
 			}
+
+			// Sleep a bit before sending download requests again to avoid a hot loop.
+			// This will only be hit if after a successful download request, there are
+			// still spans to download (e.g. because of a rabalancing).
+			time.Sleep(10 * time.Second)
 		}
 	}
 }

pkg/backup/restore_online_test.go

@@ -14,6 +14,7 @@ import (
 	"path/filepath"
 	"reflect"
 	"strings"
+	"sync/atomic"
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
@@ -32,6 +33,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/randutil"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
+	"github.com/cockroachdb/errors"
 	"github.com/kr/pretty"
 	"github.com/stretchr/testify/require"
 )
@@ -581,6 +583,65 @@ func TestOnlineRestoreErrors(t *testing.T) {
 	})
 }
 
+func TestOnlineRestoreRetryingDownloadRequests(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	defer nodelocal.ReplaceNodeLocalForTesting(t.TempDir())()
+
+	rng, seed := randutil.NewPseudoRand()
+	t.Logf("random seed: %d", seed)
+
+	alwaysFail := rng.Intn(2) == 0
+	t.Logf("always fail download requests: %t", alwaysFail)
+	totalFailures := int32(rng.Intn(maxDownloadAttempts-1) + 1)
+	var currentFailures atomic.Int32
+
+	clusterArgs := base.TestClusterArgs{
+		ServerArgs: base.TestServerArgs{
+			DefaultTestTenant: base.TestIsSpecificToStorageLayerAndNeedsASystemTenant,
+			Knobs: base.TestingKnobs{
+				BackupRestore: &sql.BackupRestoreTestingKnobs{
+					RunBeforeSendingDownloadSpan: func() error {
+						if alwaysFail {
+							return errors.Newf("always fail download request")
+						}
+						if currentFailures.Load() >= totalFailures {
+							return nil
+						}
+						currentFailures.Add(1)
+						return errors.Newf("injected download request failure")
+					},
+				},
+			},
+		},
+	}
+
+	const numAccounts = 1
+	_, sqlDB, _, cleanupFn := backupRestoreTestSetupWithParams(
+		t, singleNode, numAccounts, InitManualReplication, clusterArgs,
+	)
+	defer cleanupFn()
+
+	externalStorage := "nodelocal://1/backup"
+	sqlDB.Exec(t, fmt.Sprintf("BACKUP INTO '%s'", externalStorage))
+	sqlDB.Exec(
+		t,
+		fmt.Sprintf(`
+		RESTORE DATABASE data FROM LATEST IN '%s'
+		WITH EXPERIMENTAL DEFERRED COPY, new_db_name=data2
+		`, externalStorage),
+	)
+
+	var downloadJobID jobspb.JobID
+	sqlDB.QueryRow(t, latestDownloadJobIDQuery).Scan(&downloadJobID)
+	if alwaysFail {
+		jobutils.WaitForJobToFail(t, sqlDB, downloadJobID)
+	} else {
+		jobutils.WaitForJobToSucceed(t, sqlDB, downloadJobID)
+	}
+}
+
 func bankOnlineRestore(
 	t *testing.T,
 	sqlDB *sqlutils.SQLRunner,

pkg/server/span_download.go

@@ -18,12 +18,15 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/util/ctxgroup"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/errors/errorspb"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
 
+var perNodeMaxDownloadAttempts = 5
+
 func (t *statusServer) DownloadSpan(
 	ctx context.Context, req *serverpb.DownloadSpanRequest,
 ) (*serverpb.DownloadSpanResponse, error) {
@@ -62,8 +65,31 @@ func (s *systemStatusServer) DownloadSpan(
 	// Send DownloadSpan request to all stores on all nodes.
 	remoteRequest := *req
 	remoteRequest.NodeID = "local"
-	nodeFn := func(ctx context.Context, status serverpb.RPCStatusClient, _ roachpb.NodeID) (*serverpb.DownloadSpanResponse, error) {
-		return status.DownloadSpan(ctx, &remoteRequest)
+	nodeFn := func(
+		ctx context.Context, status serverpb.RPCStatusClient, _ roachpb.NodeID,
+	) (*serverpb.DownloadSpanResponse, error) {
+		var nodeResp *serverpb.DownloadSpanResponse
+		var err error
+		for r := retry.StartWithCtx(ctx, retry.Options{
+			InitialBackoff: time.Millisecond * 100,
+			MaxBackoff:     time.Second,
+			MaxRetries:     perNodeMaxDownloadAttempts - 1,
+		}); r.Next(); {
+			nodeResp, err = status.DownloadSpan(ctx, &remoteRequest)
+			for node, encodedErr := range resp.Errors {
+				err = errors.Wrapf(
+					errors.DecodeError(ctx, encodedErr),
+					"download span request failed with error on node n%d",
+					node,
+				)
+				break
+			}
+			if err == nil {
+				break
+			}
+			log.VInfof(ctx, 1, "attempt %d failed to download span: %v", r.CurrentAttempt(), err)
+		}
+		return nodeResp, err
 	}
 	responseFn := func(nodeID roachpb.NodeID, downloadSpanResp *serverpb.DownloadSpanResponse) {
 		for i, e := range downloadSpanResp.Errors {

pkg/sql/exec_util_backup.go

@@ -69,6 +69,10 @@ type BackupRestoreTestingKnobs struct {
 	RunAfterRestoreProcDrains func()
 
 	RunBeforeResolvingCompactionDest func() error
+
+	// RunBeforeSendingDownloadSpan is called within the retry loop of the
+	// download span worker before sending the download span request.
+	RunBeforeSendingDownloadSpan func() error
 }
 
 var _ base.ModuleTestingKnobs = &BackupRestoreTestingKnobs{}