Retry more transient errors on reads, lock release and deployment-service calls

.nextchanges/bundles/7129.md

@@ -0,0 +1 @@
+* Retry more transient server errors during `bundle deploy` and `bundle destroy`, including when releasing the deploy lock and reading state. ([#7129](https://github.com/databricks/cli/pull/7129))

bundle/deploy/filer.go

@@ -43,7 +43,9 @@ func (s stateFiler) Read(ctx context.Context, path string) (io.ReadCloser, error
 	// No get-status first: a missing file is the common first-deploy case and export's 404
 	// reports it in one round trip. A directory at a state path is not a state file; its
 	// export fails to parse as one.
-	r, err := s.workspaceClient.Workspace.Download(ctx, absPath)
+	r, err := retry.With(ctx, retry.IsTransientRead, func() (io.ReadCloser, error) {
+		return s.workspaceClient.Workspace.Download(ctx, absPath)
+	})
 	if errors.Is(err, apierr.ErrNotFound) {
 		return nil, fmt.Errorf("file does not exist: %s: %w", absPath, fs.ErrNotExist)
 	}

bundle/direct/apply.go

@@ -91,7 +91,7 @@ func (d *DeploymentUnit) Create(ctx context.Context, db *dstate.DeploymentState,
 	}
 
 	waitRemoteState, err := waitCapped(ctx, d.MaxWait, "creation of "+d.ResourceKey, func(ctx context.Context) (any, error) {
-		return retry.With(ctx, retry.IsTransient, func() (any, error) {
+		return retry.With(ctx, retry.IsTransientRead, func() (any, error) {
 			return d.Adapter.WaitAfterCreate(ctx, newID, newState)
 		})
 	})
@@ -189,7 +189,7 @@ func (d *DeploymentUnit) Update(ctx context.Context, db *dstate.DeploymentState,
 		}
 	}
 
-	waitRemoteState, err := retry.With(ctx, retry.IsTransient, func() (any, error) {
+	waitRemoteState, err := retry.With(ctx, retry.IsTransientRead, func() (any, error) {
 		return d.Adapter.WaitAfterUpdate(ctx, id, newState)
 	})
 	if err != nil {
@@ -233,7 +233,7 @@ func (d *DeploymentUnit) UpdateWithID(ctx context.Context, db *dstate.Deployment
 		return fmt.Errorf("saving state id=%s: %w", oldID, err)
 	}
 
-	waitRemoteState, err := retry.With(ctx, retry.IsTransient, func() (any, error) {
+	waitRemoteState, err := retry.With(ctx, retry.IsTransientRead, func() (any, error) {
 		return d.Adapter.WaitAfterUpdate(ctx, newID, newState)
 	})
 	if err != nil {
@@ -375,7 +375,7 @@ func (d *DeploymentUnit) refreshRemoteState(ctx context.Context, id string) erro
 	if d.RemoteState != nil {
 		return nil
 	}
-	remoteState, err := retry.With(ctx, retry.IsTransient, func() (any, error) {
+	remoteState, err := retry.With(ctx, retry.IsTransientRead, func() (any, error) {
 		return d.Adapter.DoRead(ctx, id)
 	})
 	if err != nil {

bundle/direct/bundle_plan.go

@@ -237,7 +237,7 @@ func (b *DeploymentBundle) CalculatePlan(ctx context.Context, client *databricks
 				return true
 			}
 
-			remoteState, err := retry.With(ctx, retry.IsTransient, func() (any, error) {
+			remoteState, err := retry.With(ctx, retry.IsTransientRead, func() (any, error) {
 				return adapter.DoRead(ctx, id)
 			})
 			if err != nil {
@@ -315,7 +315,7 @@ func (b *DeploymentBundle) CalculatePlan(ctx context.Context, client *databricks
 			return false
 		}
 
-		remoteState, err := retry.With(ctx, retry.IsTransient, func() (any, error) {
+		remoteState, err := retry.With(ctx, retry.IsTransientRead, func() (any, error) {
 			return adapter.DoRead(ctx, dbentry.ID)
 		})
 		if err != nil {

bundle/direct/dstate/state.go

@@ -23,6 +23,7 @@ import (
 	"github.com/databricks/cli/libs/cmdctx"
 	"github.com/databricks/cli/libs/dms"
 	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/cli/libs/retry"
 	"github.com/databricks/databricks-sdk-go/service/bundledeployments"
 	"github.com/google/uuid"
 )
@@ -584,8 +585,10 @@ Run "databricks bundle destroy" first, then deploy again with deployment history
 
 		if dmsDeployment.DeploymentID != "" {
 			w := cmdctx.WorkspaceClient(ctx)
-			recorded, err := w.BundleDeployments.ListResourcesAll(ctx, bundledeployments.ListResourcesRequest{
-				Parent: dms.DeploymentName(dmsDeployment.DeploymentID),
+			recorded, err := retry.With(ctx, retry.IsTransientRead, func() ([]bundledeployments.Resource, error) {
+				return w.BundleDeployments.ListResourcesAll(ctx, bundledeployments.ListResourcesRequest{
+					Parent: dms.DeploymentName(dmsDeployment.DeploymentID),
+				})
 			})
 			if err != nil {
 				return fmt.Errorf("listing resources from the deployment history service: %w", err)

bundle/phases/destroy.go

@@ -20,7 +20,6 @@ import (
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/logdiag"
 	"github.com/databricks/databricks-sdk-go/apierr"
-	"github.com/databricks/databricks-sdk-go/service/bundledeployments"
 )
 
 func assertRootPathExists(ctx context.Context, b *bundle.Bundle) (bool, error) {
@@ -168,7 +167,7 @@ func destroyCore(ctx context.Context, b *bundle.Bundle, plan *deployplan.Plan) {
 		if completed {
 			deploymentID := b.DeploymentBundle.StateDB.DeploymentID
 			w := b.WorkspaceClient(ctx)
-			if err := w.BundleDeployments.DeleteDeployment(ctx, bundledeployments.DeleteDeploymentRequest{Name: dms.DeploymentName(deploymentID)}); err != nil {
+			if err := deleteDeployment(ctx, w, deploymentID); err != nil {
 				logdiag.LogError(ctx, fmt.Errorf("failed to delete deployment: %w", err))
 				return
 			}
@@ -301,7 +300,7 @@ func Destroy(ctx context.Context, b *bundle.Bundle) {
 			} else if completed {
 				deploymentID := b.DeploymentBundle.StateDB.DeploymentID
 				w := b.WorkspaceClient(ctx)
-				if err := w.BundleDeployments.DeleteDeployment(ctx, bundledeployments.DeleteDeploymentRequest{Name: dms.DeploymentName(deploymentID)}); err != nil {
+				if err := deleteDeployment(ctx, w, deploymentID); err != nil {
 					logdiag.LogError(ctx, fmt.Errorf("failed to delete deployment: %w", err))
 				}
 			}

bundle/phases/dms.go

@@ -17,7 +17,10 @@ import (
 	"github.com/databricks/cli/libs/dms"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/logdiag"
+	"github.com/databricks/cli/libs/retry"
 	"github.com/databricks/cli/libs/workspaceurls"
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/common/types/fieldmask"
 	"github.com/databricks/databricks-sdk-go/service/bundledeployments"
 )
@@ -95,10 +98,12 @@ func createOrUpdateDeployment(ctx context.Context, b *bundle.Bundle, current *bu
 		}
 		db.StateDB.DeploymentID = deploymentID
 	} else if mask := metadata.StaleFields(current); mask != "" {
-		_, err := w.BundleDeployments.UpdateDeployment(ctx, bundledeployments.UpdateDeploymentRequest{
-			Name:       dms.DeploymentName(deploymentID),
-			Deployment: dms.DeploymentUpdate(metadata, mask),
-			UpdateMask: fieldmask.FieldMask{Paths: strings.Split(mask, ",")},
+		_, err := retry.With(ctx, retry.IsTransient, func() (*bundledeployments.Deployment, error) {
+			return w.BundleDeployments.UpdateDeployment(ctx, bundledeployments.UpdateDeploymentRequest{
+				Name:       dms.DeploymentName(deploymentID),
+				Deployment: dms.DeploymentUpdate(metadata, mask),
+				UpdateMask: fieldmask.FieldMask{Paths: strings.Split(mask, ",")},
+			})
 		})
 		if err != nil {
 			logdiag.LogError(ctx, fmt.Errorf("failed to update deployment: %w", err))
@@ -112,6 +117,21 @@ func createOrUpdateDeployment(ctx context.Context, b *bundle.Bundle, current *bu
 	})
 }
 
+// deleteDeployment deletes the deployment record. A retry that finds it missing succeeds: the
+// earlier attempt may have deleted it before its response was lost.
+func deleteDeployment(ctx context.Context, w *databricks.WorkspaceClient, deploymentID string) error {
+	retried := false
+	_, err := retry.With(ctx, retry.IsTransient, func() (struct{}, error) {
+		err := w.BundleDeployments.DeleteDeployment(ctx, bundledeployments.DeleteDeploymentRequest{Name: dms.DeploymentName(deploymentID)})
+		if retried && apierr.IsMissing(err) {
+			return struct{}{}, nil
+		}
+		retried = true
+		return struct{}{}, err
+	})
+	return err
+}
+
 // startVersion claims the version the run settled on and opens the buffer that records
 // each state write under it. Called after approval, so a declined deploy never claims a number.
 // A no-op when the bundle does not record deployment history.

cmd/bundle/utils/process.go

@@ -31,12 +31,14 @@ import (
 	"github.com/databricks/cli/libs/dyn"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/cli/libs/logdiag"
+	"github.com/databricks/cli/libs/retry"
 	"github.com/databricks/cli/libs/sync"
 	"github.com/databricks/cli/libs/telemetry"
 	"github.com/databricks/cli/libs/telemetry/protos"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/apierr"
 	"github.com/databricks/databricks-sdk-go/service/bundledeployments"
+	"github.com/databricks/databricks-sdk-go/service/workspace"
 	"github.com/databricks/databricks-sdk-go/useragent"
 	"github.com/spf13/cobra"
 	"golang.org/x/mod/semver"
@@ -597,16 +599,20 @@ func ValidateEngineSetting(ctx context.Context, b *bundle.Bundle) error {
 func fetchDeploymentFromStatePath(ctx context.Context, w *databricks.WorkspaceClient, statePath string) (string, *bundledeployments.Deployment, error) {
 	nodePath := path.Join(statePath, dms.DeploymentNodeName)
 
-	obj, err := w.Workspace.GetStatusByPath(ctx, nodePath)
+	obj, err := retry.With(ctx, retry.IsTransientRead, func() (*workspace.ObjectInfo, error) {
+		return w.Workspace.GetStatusByPath(ctx, nodePath)
+	})
 	if errors.Is(err, apierr.ErrNotFound) || errors.Is(err, apierr.ErrResourceDoesNotExist) {
 		return "", nil, nil
 	}
 	if err != nil {
 		return "", nil, fmt.Errorf("looking up deployment at %s: %w", nodePath, err)
 	}
 	deploymentID := strconv.FormatInt(obj.ObjectId, 10)
-	deployment, err := w.BundleDeployments.GetDeployment(ctx, bundledeployments.GetDeploymentRequest{
-		Name: dms.DeploymentName(deploymentID),
+	deployment, err := retry.With(ctx, retry.IsTransientRead, func() (*bundledeployments.Deployment, error) {
+		return w.BundleDeployments.GetDeployment(ctx, bundledeployments.GetDeploymentRequest{
+			Name: dms.DeploymentName(deploymentID),
+		})
 	})
 	if err != nil {
 		return "", nil, err

libs/locker/locker.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/databricks/cli/libs/filer"
+	"github.com/databricks/cli/libs/retry"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/google/uuid"
 )
@@ -209,7 +210,7 @@ func CreateLocker(user, targetDir string, w *databricks.WorkspaceClient) (*Locke
 	}
 
 	locker := &Locker{
-		filer: filer,
+		filer: retry.Filer(filer),
 
 		TargetDir: targetDir,
 		Active:    false,

libs/retry/filer.go

@@ -3,21 +3,24 @@ package retry
 import (
 	"bytes"
 	"context"
+	"errors"
 	"io"
+	"io/fs"
 	"slices"
 
 	"github.com/databricks/cli/libs/filer"
 )
 
-// retryingFiler retries overwriting writes that fail with a transient server error.
+// retryingFiler retries calls that fail with a transient server error and are safe to repeat.
 type retryingFiler struct {
 	filer.Filer
 }
 
-// Filer wraps f so that Write retries on transient server errors.
+// Filer wraps f so that Write, Read, ReadDir, Stat and Delete retry on transient errors.
 // Only writes with [filer.OverwriteIfExists] are retried: they are idempotent, whereas repeating
 // a write that does not overwrite could fail on the file its own earlier attempt created.
-// All other methods are passed through.
+// A retried Delete treats a missing file as success, since the earlier attempt may have removed it.
+// Mkdir is passed through.
 func Filer(f filer.Filer) filer.Filer {
 	return &retryingFiler{Filer: f}
 }
@@ -55,3 +58,34 @@ func (r *retryingFiler) Write(ctx context.Context, path string, reader io.Reader
 	})
 	return err
 }
+
+func (r *retryingFiler) Read(ctx context.Context, path string) (io.ReadCloser, error) {
+	return With(ctx, IsTransientRead, func() (io.ReadCloser, error) {
+		return r.Filer.Read(ctx, path)
+	})
+}
+
+func (r *retryingFiler) ReadDir(ctx context.Context, path string) ([]fs.DirEntry, error) {
+	return With(ctx, IsTransientRead, func() ([]fs.DirEntry, error) {
+		return r.Filer.ReadDir(ctx, path)
+	})
+}
+
+func (r *retryingFiler) Stat(ctx context.Context, path string) (fs.FileInfo, error) {
+	return With(ctx, IsTransientRead, func() (fs.FileInfo, error) {
+		return r.Filer.Stat(ctx, path)
+	})
+}
+
+func (r *retryingFiler) Delete(ctx context.Context, path string, mode ...filer.DeleteMode) error {
+	retried := false
+	_, err := With(ctx, IsTransient, func() (struct{}, error) {
+		err := r.Filer.Delete(ctx, path, mode...)
+		if retried && errors.Is(err, fs.ErrNotExist) {
+			return struct{}{}, nil
+		}
+		retried = true
+		return struct{}{}, err
+	})
+	return err
+}

libs/retry/filer_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"io"
+	"io/fs"
 	"net/http"
 	"strings"
 	"testing"
@@ -81,3 +82,87 @@ func TestFilerDoesNotRetryWithoutOverwrite(t *testing.T) {
 	require.Error(t, err)
 	assert.Len(t, f.bodies, 1)
 }
+
+// readFiler fails the first `failures` calls of every read-side method with err.
+type readFiler struct {
+	filer.Filer
+	failures int
+	err      error
+	calls    int
+}
+
+func (f *readFiler) next() error {
+	f.calls++
+	if f.calls <= f.failures {
+		return f.err
+	}
+	return nil
+}
+
+func (f *readFiler) Read(ctx context.Context, path string) (io.ReadCloser, error) {
+	return io.NopCloser(strings.NewReader("x")), f.next()
+}
+
+func (f *readFiler) ReadDir(ctx context.Context, path string) ([]fs.DirEntry, error) {
+	return nil, f.next()
+}
+
+func (f *readFiler) Stat(ctx context.Context, path string) (fs.FileInfo, error) {
+	return nil, f.next()
+}
+
+func (f *readFiler) Delete(ctx context.Context, path string, mode ...filer.DeleteMode) error {
+	return f.next()
+}
+
+func TestFilerRetriesReads(t *testing.T) {
+	ctx := testContext(t)
+	calls := map[string]func(filer.Filer) error{
+		"Read":    func(f filer.Filer) error { _, err := f.Read(ctx, "a"); return err },
+		"ReadDir": func(f filer.Filer) error { _, err := f.ReadDir(ctx, "a"); return err },
+		"Stat":    func(f filer.Filer) error { _, err := f.Stat(ctx, "a"); return err },
+	}
+	for name, call := range calls {
+		t.Run(name, func(t *testing.T) {
+			f := &readFiler{failures: 1, err: errTransient}
+			require.NoError(t, call(Filer(f)))
+			assert.Equal(t, 2, f.calls)
+		})
+	}
+}
+
+func TestFilerDoesNotRetryNotExist(t *testing.T) {
+	f := &readFiler{failures: 1, err: fs.ErrNotExist}
+	_, err := Filer(f).Stat(testContext(t), "a")
+	require.ErrorIs(t, err, fs.ErrNotExist)
+	assert.Equal(t, 1, f.calls)
+}
+
+func TestFilerDeleteRetries(t *testing.T) {
+	f := &readFiler{failures: 1, err: errTransient}
+	require.NoError(t, Filer(f).Delete(testContext(t), "a"))
+	assert.Equal(t, 2, f.calls)
+}
+
+func TestFilerDeleteMissingAfterRetryIsSuccess(t *testing.T) {
+	// The first attempt hits a transient error, the retry finds the file gone.
+	f := &sequenceFiler{errs: []error{errTransient, fs.ErrNotExist}}
+	require.NoError(t, Filer(f).Delete(testContext(t), "a"))
+}
+
+func TestFilerDeleteMissingOnFirstAttemptIsError(t *testing.T) {
+	f := &sequenceFiler{errs: []error{fs.ErrNotExist}}
+	require.ErrorIs(t, Filer(f).Delete(testContext(t), "a"), fs.ErrNotExist)
+}
+
+// sequenceFiler returns errs from successive Delete calls.
+type sequenceFiler struct {
+	filer.Filer
+	errs []error
+}
+
+func (f *sequenceFiler) Delete(ctx context.Context, path string, mode ...filer.DeleteMode) error {
+	err := f.errs[0]
+	f.errs = f.errs[1:]
+	return err
+}