From 4288b8132eb1f03d6b9bdb8f0cc0dd8d5f9c8a04 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Wed, 21 Oct 2026 16:13:39 +0000 Subject: [PATCH 01/52] Add explicit deployment state modes and optional DMS history --- bundle/direct/bind.go | 12 +- bundle/direct/dstate/mode.go | 45 +++ bundle/direct/dstate/mode_test.go | 421 +++++++++++++++++++++++++++ bundle/direct/dstate/state.go | 307 ++++++++++--------- bundle/direct/dstate/state_test.go | 71 +++-- bundle/migrate/build_state_test.go | 2 +- bundle/phases/deploy.go | 7 +- bundle/phases/destroy.go | 7 +- bundle/phases/dms.go | 2 +- bundle/statemgmt/direct_migration.go | 4 +- cmd/bundle/deployment/migrate.go | 2 +- cmd/bundle/utils/process.go | 8 +- libs/dms/operation_buffer.go | 74 ++++- libs/dms/operation_buffer_test.go | 36 ++- 14 files changed, 807 insertions(+), 191 deletions(-) create mode 100644 bundle/direct/dstate/mode.go create mode 100644 bundle/direct/dstate/mode_test.go diff --git a/bundle/direct/bind.go b/bundle/direct/bind.go index 9ac921b9c3c..7b4aca63ee6 100644 --- a/bundle/direct/bind.go +++ b/bundle/direct/bind.go @@ -66,7 +66,7 @@ func (b *DeploymentBundle) Bind(ctx context.Context, client *databricks.Workspac // The state is opened without a DMS client, so the writes below record nothing; // phases.Bind and phases.Unbind refuse to run at all when recording is enabled. var checkStateDB dstate.DeploymentState - if err := checkStateDB.Open(ctx, statePath, dstate.WithRecovery(true), dstate.WithWrite(false), dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{}); err != nil { + if err := checkStateDB.Open(ctx, statePath, dstate.ReadingFile, dstate.WithRecovery(true), dstate.OpenDmsArgs{}); err != nil { // State that cannot be read is not the same as state without a binding: // the resource may well be managed already, and binding on top of it // would take over a resource whose ownership was never checked. A state @@ -95,7 +95,7 @@ func (b *DeploymentBundle) Bind(ctx context.Context, client *databricks.Workspac } // Open temp state - err := b.StateDB.Open(ctx, tmpStatePath, dstate.WithRecovery(false), dstate.WithWrite(true), dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{}) + err := b.StateDB.Open(ctx, tmpStatePath, dstate.WritingFile, dstate.WithRecovery(false), dstate.OpenDmsArgs{}) if err != nil { os.Remove(tmpStatePath) return nil, err @@ -118,7 +118,7 @@ func (b *DeploymentBundle) Bind(ctx context.Context, client *databricks.Workspac log.Infof(ctx, "Bound %s to id=%s (in temp state)", resourceKey, resourceID) // First plan + update: populate state with resolved config - err = b.StateDB.Open(ctx, tmpStatePath, dstate.WithRecovery(true), dstate.WithWrite(false), dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{}) + err = b.StateDB.Open(ctx, tmpStatePath, dstate.ReadingFile, dstate.WithRecovery(true), dstate.OpenDmsArgs{}) if err != nil { os.Remove(tmpStatePath) return nil, err @@ -167,7 +167,7 @@ func (b *DeploymentBundle) Bind(ctx context.Context, client *databricks.Workspac return nil, fmt.Errorf("compacting state: %w", err) } - err = b.StateDB.Open(ctx, tmpStatePath, dstate.WithRecovery(true), dstate.WithWrite(true), dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{}) + err = b.StateDB.Open(ctx, tmpStatePath, dstate.WritingFile, dstate.WithRecovery(true), dstate.OpenDmsArgs{}) if err != nil { os.Remove(tmpStatePath) return nil, err @@ -187,7 +187,7 @@ func (b *DeploymentBundle) Bind(ctx context.Context, client *databricks.Workspac } // Second plan: this is the plan to present to the user (change between remote resource and config) - err = b.StateDB.Open(ctx, tmpStatePath, dstate.WithRecovery(true), dstate.WithWrite(false), dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{}) + err = b.StateDB.Open(ctx, tmpStatePath, dstate.ReadingFile, dstate.WithRecovery(true), dstate.OpenDmsArgs{}) if err != nil { os.Remove(tmpStatePath) return nil, err @@ -237,7 +237,7 @@ func (result *BindResult) Cancel() { // Unbind removes a resource from direct engine state without deleting // the workspace resource. Also removes associated permissions/grants entries. func (b *DeploymentBundle) Unbind(ctx context.Context, statePath, resourceKey string) error { - err := b.StateDB.Open(ctx, statePath, dstate.WithRecovery(true), dstate.WithWrite(true), dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{}) + err := b.StateDB.Open(ctx, statePath, dstate.WritingFile, dstate.WithRecovery(true), dstate.OpenDmsArgs{}) if err != nil { if errors.Is(err, dstate.ErrUnsettingRecording) { return errors.New("unbind is not supported for a bundle target that records deployment history") diff --git a/bundle/direct/dstate/mode.go b/bundle/direct/dstate/mode.go new file mode 100644 index 00000000000..07fa8f61228 --- /dev/null +++ b/bundle/direct/dstate/mode.go @@ -0,0 +1,45 @@ +package dstate + +import "fmt" + +// Mode selects the source of resource state and the persistence required for writes. +type Mode int + +const ( + ReadingFile Mode = iota + WritingFile + ReadingDMS + WritingDMS + WritingFileWithHistory + WritingFileWithHistoryEnforced +) + +func (m Mode) isWritable() bool { + switch m { + case WritingFile, WritingDMS, WritingFileWithHistory, WritingFileWithHistoryEnforced: + return true + default: + return false + } +} + +func (m Mode) usesDMSState() bool { + return m == ReadingDMS || m == WritingDMS +} + +func (m Mode) isFileWithHistory() bool { + return m == WritingFileWithHistory || m == WritingFileWithHistoryEnforced +} + +func (m Mode) recordsHistory() bool { + return m.usesDMSState() || m.isFileWithHistory() +} + +func (m Mode) validate() error { + switch m { + case ReadingFile, WritingFile, ReadingDMS, WritingDMS, WritingFileWithHistory, WritingFileWithHistoryEnforced: + return nil + default: + return fmt.Errorf("invalid deployment state mode: %d", m) + } +} diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go new file mode 100644 index 00000000000..73d5e348a42 --- /dev/null +++ b/bundle/direct/dstate/mode_test.go @@ -0,0 +1,421 @@ +package dstate_test + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/databricks/cli/bundle/deployplan" + "github.com/databricks/cli/bundle/direct/dstate" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/databricks-sdk-go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type modeTransport struct { + handler http.HandlerFunc +} + +func (m modeTransport) RoundTrip(r *http.Request) (*http.Response, error) { + w := httptest.NewRecorder() + if r.URL.Path == "/.well-known/databricks-config" || m.handler == nil { + w.WriteHeader(http.StatusNotFound) + } else { + m.handler.ServeHTTP(w, r) + } + response := w.Result() + response.Request = r + return response, nil +} + +// modeContext supplies an in-memory DMS endpoint without resolving workspace authentication. +func modeContext(t *testing.T, handler http.HandlerFunc) context.Context { + t.Helper() + client, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: "https://state.test", Token: "test-token", AuthType: "pat", + HTTPTransport: modeTransport{handler: handler}, + }) + require.NoError(t, err) + return cmdctx.SetWorkspaceClient(t.Context(), client) +} + +func TestModeAssertions(t *testing.T) { + for _, tc := range []struct { + mode dstate.Mode + writable bool + dms bool + history bool + }{ + {dstate.ReadingFile, false, false, false}, + {dstate.WritingFile, true, false, false}, + {dstate.ReadingDMS, false, true, true}, + {dstate.WritingDMS, true, true, true}, + {dstate.WritingFileWithHistory, true, false, true}, + {dstate.WritingFileWithHistoryEnforced, true, false, true}, + } { + t.Run(fmt.Sprint(tc.mode), func(t *testing.T) { + ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + }) + path := filepath.Join(t.TempDir(), "resources.json") + var db dstate.DeploymentState + require.NoError(t, db.Open(ctx, path, tc.mode, false, dstate.OpenDmsArgs{})) + assert.Equal(t, tc.dms, db.IsDeploymentMetadataService()) + assert.Equal(t, tc.history, db.RecordsHistory()) + assert.NotPanics(t, db.AssertOpenedForReadOrWrite) + if tc.writable { + assert.NotPanics(t, db.AssertOpenedForWrite) + assert.Panics(t, db.AssertOpenedForRead) + } else { + assert.NotPanics(t, db.AssertOpenedForRead) + assert.Panics(t, db.AssertOpenedForWrite) + } + if tc.writable && !tc.dms { + assert.FileExists(t, path+".wal") + } else { + assert.NoFileExists(t, path+".wal") + } + _, err := db.Finalize(ctx) + require.NoError(t, err) + assert.Panics(t, db.AssertOpenedForReadOrWrite) + assert.NoFileExists(t, path+".wal") + }) + } +} + +func TestFileModesPersistStateAndHistory(t *testing.T) { + for _, mode := range []dstate.Mode{dstate.WritingFile, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { + t.Run(fmt.Sprint(mode), func(t *testing.T) { + var mu sync.Mutex + updates := make(map[string]map[string]any) + ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodPatch, r.Method) + var body map[string]any + if !assert.NoError(t, json.NewDecoder(r.Body).Decode(&body)) { + w.WriteHeader(http.StatusBadRequest) + return + } + mu.Lock() + updates[filepath.Base(r.URL.Path)] = body + mu.Unlock() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"sequence_id":1}`)) + }) + path := filepath.Join(t.TempDir(), "resources.json") + var db dstate.DeploymentState + require.NoError(t, db.Open(ctx, path, mode, true, dstate.OpenDmsArgs{})) + if mode != dstate.WritingFile { + db.InitializeOperationBuffer(ctx, "dep-1", 1, 2) + } + dependsOn := []deployplan.DependsOnEntry{{Node: "resources.pipelines.upstream", Label: "upstream"}} + require.NoError(t, db.SaveState(ctx, "resources.jobs.kept", "keep-id", map[string]string{"name": "kept"}, dependsOn)) + require.NoError(t, db.SaveState(ctx, "resources.jobs.deleted", "delete-id", map[string]string{"name": "deleted"}, nil)) + require.NoError(t, db.DeleteState(ctx, "resources.jobs.deleted", false)) + exported, err := db.Finalize(ctx) + require.NoError(t, err) + assert.Len(t, exported, 1) + assert.Equal(t, "keep-id", exported["resources.jobs.kept"].ID) + assert.NoFileExists(t, path+".wal") + + var reopened dstate.DeploymentState + require.NoError(t, reopened.Open(t.Context(), path, dstate.ReadingFile, false, dstate.OpenDmsArgs{})) + assert.Equal(t, 1, reopened.GetSerial()) + assert.Len(t, reopened.Data.State, 1) + assert.Equal(t, dependsOn, reopened.Data.State["resources.jobs.kept"].DependsOn) + _, enforced := reopened.Data.Features[dstate.FeatureDMSShadow] + assert.Equal(t, mode == dstate.WritingFileWithHistoryEnforced, enforced) + _, err = reopened.Finalize(t.Context()) + require.NoError(t, err) + + mu.Lock() + defer mu.Unlock() + if mode == dstate.WritingFile { + assert.Empty(t, updates) + } else { + require.Len(t, updates, 2) + assert.JSONEq(t, `{"state":{"name":"kept"},"depends_on":[{"node":"resources.pipelines.upstream","label":"upstream"}]}`, updates["jobs.kept"]["state"].(string)) + assert.NotContains(t, updates["jobs.deleted"], "state") + } + }) + } +} + +func TestDMSReadsAndWritesWithoutWAL(t *testing.T) { + ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodGet { + _, _ = w.Write([]byte(`{"resources":[{"resource_key":"jobs.remote","resource_id":"remote-id","state":"{\"state\":{\"name\":\"remote\"}}"}]}`)) + } else { + _, _ = w.Write([]byte(`{"sequence_id":1}`)) + } + }) + path := filepath.Join(t.TempDir(), "resources.json") + var db dstate.DeploymentState + require.NoError(t, db.Open(ctx, path, dstate.ReadingDMS, false, dstate.OpenDmsArgs{DeploymentID: "dep-1", LastVersionID: 7})) + assert.Equal(t, "remote-id", db.GetResourceID("resources.jobs.remote")) + assert.Equal(t, 7, db.GetSerial()) + require.NoError(t, db.UpgradeToWrite(dstate.WritingDMS)) + assert.NoFileExists(t, path+".wal") + db.InitializeOperationBuffer(ctx, "dep-1", 8, 2) + require.NoError(t, db.SaveState(ctx, "resources.jobs.new", "new-id", map[string]string{"name": "new"}, nil)) + require.NoError(t, db.DeleteState(ctx, "resources.jobs.remote", false)) + exported, err := db.Finalize(ctx) + require.NoError(t, err) + assert.Len(t, exported, 1) + assert.Equal(t, "new-id", exported["resources.jobs.new"].ID) + assert.NoFileExists(t, path+".wal") + data, err := os.ReadFile(path) + require.NoError(t, err) + var marker dstate.Database + require.NoError(t, json.Unmarshal(data, &marker)) + assert.Empty(t, marker.State) + assert.Equal(t, 8, marker.Serial) + assert.Contains(t, marker.Features, dstate.FeatureDeploymentHistory) +} + +func TestUpgradeToWriteValidatesMode(t *testing.T) { + for _, tc := range []struct { + name string + read dstate.Mode + write dstate.Mode + shadow bool + wantErr string + }{ + {"file", dstate.ReadingFile, dstate.WritingFile, false, ""}, + {"history", dstate.ReadingFile, dstate.WritingFileWithHistory, false, ""}, + {"enforced history", dstate.ReadingFile, dstate.WritingFileWithHistoryEnforced, false, ""}, + {"DMS", dstate.ReadingDMS, dstate.WritingDMS, false, ""}, + {"file to DMS", dstate.ReadingFile, dstate.WritingDMS, false, "cannot change"}, + {"DMS to file", dstate.ReadingDMS, dstate.WritingFile, false, "cannot change"}, + {"DMS to history", dstate.ReadingDMS, dstate.WritingFileWithHistory, false, "cannot change"}, + {"reading mode", dstate.ReadingFile, dstate.ReadingFile, false, "invalid"}, + {"unknown mode", dstate.ReadingFile, dstate.Mode(99), false, "invalid"}, + {"enforcement retained", dstate.ReadingFile, dstate.WritingFileWithHistory, true, ""}, + {"enforcement bypass", dstate.ReadingFile, dstate.WritingFile, true, "dms_shadow"}, + } { + t.Run(tc.name, func(t *testing.T) { + ctx := modeContext(t, nil) + path := filepath.Join(t.TempDir(), "resources.json") + if tc.shadow { + require.NoError(t, os.WriteFile(path, []byte(`{"state_version":3,"features":{"dms_shadow":{}},"state":{}}`), 0o600)) + } + var db dstate.DeploymentState + require.NoError(t, db.Open(ctx, path, tc.read, false, dstate.OpenDmsArgs{})) + err := db.UpgradeToWrite(tc.write) + if tc.wantErr != "" { + require.ErrorContains(t, err, tc.wantErr) + assert.NotPanics(t, db.AssertOpenedForRead) + assert.NoFileExists(t, path+".wal") + } else { + require.NoError(t, err) + assert.NotPanics(t, db.AssertOpenedForWrite) + } + _, err = db.Finalize(ctx) + require.NoError(t, err) + }) + } +} + +func TestHistoryErrorsFailOpen(t *testing.T) { + for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { + t.Run(fmt.Sprint(mode), func(t *testing.T) { + ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error_code":"INVALID_PARAMETER_VALUE","message":"history unavailable"}`)) + }) + path := filepath.Join(t.TempDir(), "resources.json") + var db dstate.DeploymentState + require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) + db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) + require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", map[string]string{"name": "job"}, nil)) + _, err := db.Finalize(ctx) + if mode == dstate.WritingDMS { + require.ErrorContains(t, err, "history unavailable") + assert.Error(t, db.RecordingError()) + } else { + require.NoError(t, err) + assert.NoError(t, db.RecordingError()) + var reopened dstate.DeploymentState + require.NoError(t, reopened.Open(t.Context(), path, dstate.ReadingFile, false, dstate.OpenDmsArgs{})) + assert.Equal(t, "id-1", reopened.GetResourceID("resources.jobs.job")) + _, err = reopened.Finalize(t.Context()) + require.NoError(t, err) + } + _, err = db.CompleteVersion(ctx, true) + if mode == dstate.WritingDMS { + assert.ErrorContains(t, err, "history unavailable") + } else { + assert.NoError(t, err) + } + completed, err := db.CompleteVersion(ctx, true) + require.NoError(t, err) + assert.False(t, completed) + }) + } +} + +func TestHistoryDrainTimeoutPreservesFileState(t *testing.T) { + for _, mode := range []dstate.Mode{dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { + t.Run(fmt.Sprint(mode), func(t *testing.T) { + started := make(chan struct{}) + finished := make(chan struct{}) + ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + close(started) + <-r.Context().Done() + close(finished) + }) + path := filepath.Join(t.TempDir(), "resources.json") + var db dstate.DeploymentState + require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) + db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) + require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) + select { + case <-started: + case <-time.After(5 * time.Second): + t.Fatal("DMS request did not start") + } + drainCtx, cancel := context.WithTimeout(ctx, 20*time.Millisecond) + defer cancel() + _, err := db.Finalize(drainCtx) + require.NoError(t, err) + assert.FileExists(t, path) + assert.NoFileExists(t, path+".wal") + select { + case <-finished: + case <-time.After(5 * time.Second): + t.Fatal("DMS request was not canceled") + } + }) + } +} + +func TestEnforcedHistorySurvivesWALRecovery(t *testing.T) { + for _, withResource := range []bool{false, true} { + t.Run(map[bool]string{false: "feature only", true: "with resource"}[withResource], func(t *testing.T) { + ctx := modeContext(t, nil) + path := filepath.Join(t.TempDir(), "resources.json") + var db dstate.DeploymentState + require.NoError(t, db.Open(ctx, path, dstate.WritingFileWithHistoryEnforced, false, dstate.OpenDmsArgs{})) + if withResource { + require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) + } + wal, err := os.ReadFile(path + ".wal") + require.NoError(t, err) + _, err = db.Finalize(ctx) + require.NoError(t, err) + // Reproduce an interrupted deployment using only its WAL. + recoveredPath := filepath.Join(t.TempDir(), "resources.json") + require.NoError(t, os.WriteFile(recoveredPath+".wal", wal, 0o600)) + var recovered dstate.DeploymentState + require.NoError(t, recovered.Open(t.Context(), recoveredPath, dstate.ReadingFile, true, dstate.OpenDmsArgs{})) + assert.Contains(t, recovered.StateFeatures(), dstate.FeatureDMSShadow) + assert.Equal(t, 1, recovered.GetSerial()) + assert.FileExists(t, recoveredPath) + assert.NoFileExists(t, recoveredPath+".wal") + require.ErrorContains(t, recovered.UpgradeToWrite(dstate.WritingFile), "dms_shadow") + if withResource { + assert.Equal(t, "id-1", recovered.GetResourceID("resources.jobs.job")) + } + _, err = recovered.Finalize(ctx) + require.NoError(t, err) + }) + } +} + +func TestOpenRejectsInvalidModesAndFeatures(t *testing.T) { + for _, tc := range []struct { + name string + mode dstate.Mode + features string + wantErr string + }{ + {"unknown mode", dstate.Mode(99), `{}`, "invalid deployment state mode"}, + {"unknown feature", dstate.ReadingFile, `{"future":{}}`, "requires features"}, + {"conflicting features", dstate.ReadingFile, `{"deployment_history":{},"dms_shadow":{}}`, "dms_shadow"}, + {"DMS reading shadow", dstate.ReadingDMS, `{"dms_shadow":{}}`, "dms_shadow"}, + {"file writing DMS", dstate.WritingFile, `{"deployment_history":{}}`, "not supported"}, + {"file writing shadow", dstate.WritingFile, `{"dms_shadow":{}}`, "dms_shadow"}, + } { + t.Run(tc.name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "resources.json") + seed := `{"state_version":3,"features":` + tc.features + `,"state":{}}` + require.NoError(t, os.WriteFile(path, []byte(seed), 0o600)) + var db dstate.DeploymentState + err := db.Open(t.Context(), path, tc.mode, false, dstate.OpenDmsArgs{}) + require.ErrorContains(t, err, tc.wantErr) + assert.False(t, db.IsOpen()) + assert.NoFileExists(t, path+".wal") + }) + } +} + +func TestCompleteVersionAfterFinalize(t *testing.T) { + for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { + t.Run(fmt.Sprint(mode), func(t *testing.T) { + completions := 0 + ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodPost { + completions++ + assert.Equal(t, "/api/2.0/bundle/deployments/dep-1/versions/1/complete", r.URL.Path) + var body map[string]string + assert.NoError(t, json.NewDecoder(r.Body).Decode(&body)) + assert.Equal(t, "VERSION_COMPLETE_SUCCESS", body["completion_reason"]) + _, _ = w.Write([]byte(`{}`)) + } else { + _, _ = w.Write([]byte(`{"sequence_id":1}`)) + } + }) + path := filepath.Join(t.TempDir(), "resources.json") + var db dstate.DeploymentState + require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) + db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) + require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) + _, err := db.Finalize(ctx) + require.NoError(t, err) + assert.False(t, db.IsOpen()) + completed, err := db.CompleteVersion(ctx, true) + require.NoError(t, err) + assert.True(t, completed) + completed, err = db.CompleteVersion(ctx, true) + require.NoError(t, err) + assert.False(t, completed) + assert.Equal(t, 1, completions) + }) + } +} + +func TestFailedUpgradeLeavesReadModeAvailable(t *testing.T) { + path := filepath.Join(t.TempDir(), "resources.json") + var db dstate.DeploymentState + require.NoError(t, db.Open(t.Context(), path, dstate.ReadingFile, true, dstate.OpenDmsArgs{})) + require.NoError(t, os.WriteFile(path+".wal", []byte("occupied"), 0o600)) + require.Error(t, db.UpgradeToWrite(dstate.WritingFile)) + assert.NotPanics(t, db.AssertOpenedForRead) + require.NoError(t, os.Remove(path+".wal")) + require.NoError(t, db.UpgradeToWrite(dstate.WritingFile)) + _, err := db.Finalize(t.Context()) + require.NoError(t, err) +} + +func TestRecoveryRejectsUnknownWALFeatures(t *testing.T) { + path := filepath.Join(t.TempDir(), "resources.json") + wal := `{"state_version":3,"lineage":"lineage","serial":1,"features":{"future":{}}}` + "\n" + `{"k":"resources.jobs.job","v":{"__id__":"id-1","state":{}}}` + "\n" + require.NoError(t, os.WriteFile(path+".wal", []byte(wal), 0o600)) + var db dstate.DeploymentState + require.ErrorContains(t, db.Open(t.Context(), path, dstate.ReadingFile, true, dstate.OpenDmsArgs{}), "future") + assert.False(t, db.IsOpen()) + assert.NoFileExists(t, path) + assert.FileExists(t, path+".wal") +} diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 4d2a82546f9..8666f7c8cd0 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -14,6 +14,7 @@ import ( "slices" "strings" "sync" + "time" "uuid" "github.com/databricks/cli/bundle/config" @@ -31,10 +32,11 @@ import ( const ( // currentStateVersion is the schema version written by this CLI and the version // legacy states are migrated up to on load. - currentStateVersion = 3 - initialBufferSize = 64 * 1024 - maxWalEntrySize = 10 * 1024 * 1024 - walSuffix = ".wal" + currentStateVersion = 3 + initialBufferSize = 64 * 1024 + maxWalEntrySize = 10 * 1024 * 1024 + walSuffix = ".wal" + historyOperationTimeout = 5 * time.Second ) // FeatureDeploymentHistory marks a state whose resources are also recorded with the @@ -47,6 +49,9 @@ const ( // deploying such a state without recording is refused. const FeatureDeploymentHistory = "deployment_history" +// FeatureDMSShadow requires writers to record history while keeping file state authoritative. +const FeatureDMSShadow = "dms_shadow" + // featuresDocURL explains the state-features mechanism in the error shown when this CLI refuses a // state that depends on a feature it does not recognize. const featuresDocURL = "https://docs.databricks.com/aws/en/dev-tools/bundles/state-features#state-features" @@ -56,6 +61,7 @@ const featuresDocURL = "https://docs.databricks.com/aws/en/dev-tools/bundles/sta // silently clobbered by this one. var recognizedFeatures = map[string]struct{}{ FeatureDeploymentHistory: {}, + FeatureDMSShadow: {}, } // assertNoUnsupportedFeatures refuses a state that depends on a feature this CLI does not recognize. @@ -88,9 +94,8 @@ type DeploymentState struct { mu sync.Mutex walFile *os.File - // openedForWrite records write mode explicitly. It used to be implied by walFile, but a recorded - // deployment writes no WAL, so the two are no longer the same question. - openedForWrite bool + // mode survives Finalize so version completion retains its error policy. + mode Mode // Maps resource key to ID. Unlike Data.State, this is up to date during writes (deploys). stateIDs map[string]string @@ -103,10 +108,7 @@ type DeploymentState struct { // completion after an explicit one does nothing. versionCompleted bool - // recordsHistory is whether this deployment records history, decided by Open from the state's - // features. Cached because it outlives the features themselves: Finalize resets Data, and - // CompleteVersion runs after that. - recordsHistory bool + historyDrainTimeout time.Duration // DeploymentID is the recorded deployment's id. Set at Open, or by the first deploy that creates // the deployment, which Open cannot know about. @@ -177,12 +179,17 @@ type WALEntry struct { // Open built, so it is called only when the deployment records history, and once the version exists // (after approval) - which is why it is not an Open option. It also records the id a first deploy // just created, which Open could not know, so CompleteVersion later has it. -func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploymentID string, versionID int) { - buf := dms.StartOperationBuffer(ctx, deploymentID, versionID) - +func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploymentID string, versionID, operationCount int) { + db.AssertOpenedForWrite() db.mu.Lock() defer db.mu.Unlock() - db.operationBuffer = buf + options := dms.BufferOptions{} + if db.mode.isFileWithHistory() { + options.BestEffort = true + options.Capacity = operationCount + db.historyDrainTimeout = time.Duration(operationCount) * historyOperationTimeout + } + db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options) db.DeploymentID = deploymentID db.VersionID = versionID } @@ -195,9 +202,12 @@ func (db *DeploymentState) getOperationBuffer() *dms.OperationBuffer { } // RecordingError reports whether recording state writes to the service has failed, so the apply -// stops touching resources once the service is no longer keeping up. Nil when the bundle does not -// record deployment history or recording is healthy. +// stops touching resources once the service is no longer keeping up. File-backed history modes +// always return nil because their resource state remains recoverable from the WAL. func (db *DeploymentState) RecordingError() error { + if !db.IsDeploymentMetadataService() { + return nil + } buf := db.getOperationBuffer() if buf == nil { return nil @@ -211,10 +221,6 @@ func (db *DeploymentState) RecordingError() error { // one does nothing. Finalize has already drained the buffered operations. It reads only fields that // survive Finalize's reset and asserts nothing, so it is safe to call after the state is closed. func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (bool, error) { - if !db.IsDeploymentMetadataService() { - return false, nil - } - db.mu.Lock() buf := db.operationBuffer if buf == nil || db.versionCompleted { @@ -223,8 +229,15 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b } db.versionCompleted = true deploymentID, versionID := db.DeploymentID, db.VersionID + bestEffort := db.mode.isFileWithHistory() db.mu.Unlock() + if bestEffort { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, historyOperationTimeout) + defer cancel() + } + // A recording failure fails the version even when the caller counted the deploy a success: the // service does not then hold everything the WAL does. Finalize already drained and surfaced it; // this reads the drained buffer's error so a destroy whose uploads failed keeps its record. @@ -242,6 +255,10 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b CompletionReason: reason, }) if err != nil { + if bestEffort { + log.Warnf(ctx, "Failed to complete deployment history version: %s", err) + return false, nil + } return false, err } log.Infof(ctx, "Completed deployment version: deployment=%s version=%d reason=%s", deploymentID, versionID, reason) @@ -253,7 +270,7 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b // RecordFailure records that a resource did not apply, so the history says why rather than // leaving the resource out. resourceID is the id it had before the failure. func (db *DeploymentState) RecordFailure(resourceKey, resourceID string, cause error) { - if !db.IsDeploymentMetadataService() { + if !db.RecordsHistory() { return } @@ -307,27 +324,22 @@ func (db *DeploymentState) SaveState(ctx context.Context, key, newID string, sta DependsOn: dependsOn, } - // A recorded deployment persists through the service, everything else through the WAL. The - // entry is still kept in memory: Finalize exports it for metadata.json and the deploy summary, - // and dataForFile empties State again before the tombstone is written. - if db.isDeploymentMetadataService() { + if db.mode.usesDMSState() { db.Data.State[key] = entry - db.stateIDs[key] = newID - if buf := db.operationBuffer; buf != nil { - recorded, err := json.Marshal(RecordedState{State: entry.State, DependsOn: dependsOn}) - if err != nil { - return err - } - buf.RecordOperation(ctx, key, false, newID, recorded) - } - return nil + } else if err := appendJSONLine(db.walFile, WALEntry{Key: key, Value: &entry}); err != nil { + return err } + db.stateIDs[key] = newID - err = appendJSONLine(db.walFile, WALEntry{Key: key, Value: &entry}) - if err == nil { - db.stateIDs[key] = newID + // File-backed modes record history only after the WAL append succeeds. + if buf := db.operationBuffer; buf != nil { + recorded, err := json.Marshal(RecordedState{State: entry.State, DependsOn: dependsOn}) + if err != nil { + return err + } + buf.RecordOperation(ctx, key, false, newID, recorded) } - return err + return nil } // DeleteState drops the resource's state entry: the resource is gone. inProgress records the @@ -344,14 +356,14 @@ func (db *DeploymentState) DeleteState(ctx context.Context, key string, inProgre // Read before the delete below: DMS needs the id to say which resource went away. deletedID := db.stateIDs[key] - // A recorded deployment persists through the service, everything else through the WAL. - if db.isDeploymentMetadataService() { - if buf := db.operationBuffer; buf != nil { - buf.RecordOperation(ctx, key, inProgress, deletedID, nil) - } + if db.mode.usesDMSState() { + delete(db.Data.State, key) } else if err := appendJSONLine(db.walFile, WALEntry{Key: key}); err != nil { return err } + if buf := db.operationBuffer; buf != nil { + buf.RecordOperation(ctx, key, inProgress, deletedID, nil) + } delete(db.stateIDs, key) return nil @@ -415,7 +427,7 @@ func (db *DeploymentState) GetSerial() int { // isDeploymentMetadataService is IsDeploymentMetadataService for callers already holding db.mu. func (db *DeploymentState) isDeploymentMetadataService() bool { - return db.recordsHistory + return db.mode.usesDMSState() } // IsDeploymentMetadataService reports whether this deployment's resource state lives in the @@ -427,6 +439,13 @@ func (db *DeploymentState) IsDeploymentMetadataService() bool { return db.isDeploymentMetadataService() } +// RecordsHistory reports whether this mode records resource operations with DMS. +func (db *DeploymentState) RecordsHistory() bool { + db.mu.Lock() + defer db.mu.Unlock() + return db.mode.recordsHistory() +} + // GetOrInitLineage returns the deployment lineage, generating and storing a new // one if the state does not have one yet. It is the single place the lineage is // initialized, shared so the direct deployment engine (when it writes state, via @@ -447,27 +466,12 @@ func (db *DeploymentState) GetOrInitLineage() string { return db.Data.Lineage } -type ( - // If true, then Open reads the WAL and merges it in the state. If false, and WAL is present, Open returns an error. - WithRecovery bool - - // If true, the state is opened in Write mode, which enables methods such as SaveState - // but disables GetResourceEntry (since writes go strictly into WAL and not in memory). - WithWrite bool - - // If true, the deployment records history with the metadata service: Open builds a DMS - // client from the workspace client, reads resources from the service, and refuses a state - // that tracks resources without the recording marker. It forces WithRecovery off, since the - // service is the source of truth and a leftover WAL is discarded rather than replayed. - WithDeploymentHistory bool -) +// WithRecovery allows Open to merge a leftover WAL into file state. +type WithRecovery bool -// Open reads the deployment state from disk, recovering the WAL when withRecovery is set. -// When withDeploymentHistory is set it builds a DMS client from wsClient and reads resources from -// the service instead, with dmsDeploymentID the id the service holds (empty before the first -// recorded deploy). Lineage comes from the file. Open only reads through the client - -// InitializeOperationBuffer installs the write path once a version exists. -func (db *DeploymentState) Open(ctx context.Context, path string, withRecovery WithRecovery, withWrite WithWrite, withDeploymentHistory WithDeploymentHistory, dmsDeployment OpenDmsArgs) error { +// Open loads resource state using mode, recovering a leftover WAL when requested. +// DMS reads and history writes require a workspace client on ctx. +func (db *DeploymentState) Open(ctx context.Context, path string, mode Mode, withRecovery WithRecovery, dmsDeployment OpenDmsArgs) error { db.mu.Lock() defer db.mu.Unlock() @@ -475,13 +479,12 @@ func (db *DeploymentState) Open(ctx context.Context, path string, withRecovery W panic(fmt.Sprintf("state already opened: %v, cannot open %v", db.Path, path)) } - err := db.unlockedOpen(ctx, path, withRecovery, withWrite, withDeploymentHistory, dmsDeployment) + err := db.unlockedOpen(ctx, path, mode, withRecovery, dmsDeployment) if err != nil { - // A failed open must leave the receiver closed. unlockedOpen assigns - // db.Path before every fallible step, so without this the receiver stays - // half-initialized and the next Open on it hits the panic above instead - // of reporting the real error. + // A failed open must leave the receiver closed so it can be retried. db.reset() + db.mode = ReadingFile + db.operationBuffer = nil } return err } @@ -495,14 +498,22 @@ func (db *DeploymentState) reset() { db.Path = "" db.Data = Database{} db.stateIDs = nil - db.openedForWrite = false } -func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, withRecovery WithRecovery, withWrite WithWrite, withDeploymentHistory WithDeploymentHistory, dmsDeployment OpenDmsArgs) error { +func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode Mode, withRecovery WithRecovery, dmsDeployment OpenDmsArgs) error { + if err := mode.validate(); err != nil { + return err + } db.Path = path - // Cleared here rather than in reset, which CompleteVersion needs it to survive: the same state - // is reopened (see bind.go), and a stale value would misroute a non-recording open. - db.recordsHistory = false + db.mode = ReadingFile + if mode.usesDMSState() { + db.mode = ReadingDMS + } + db.operationBuffer = nil + db.versionCompleted = false + db.historyDrainTimeout = 0 + db.DeploymentID = dmsDeployment.DeploymentID + db.VersionID = dmsDeployment.LastVersionID // The state file is the source of truth for whether this deployment records history: read it // first. @@ -528,7 +539,7 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, withRe db.stateIDs[key] = entry.ID } - recording := bool(withDeploymentHistory) + recording := mode.usesDMSState() walPath := db.Path + walSuffix _, err = os.Stat(walPath) @@ -561,6 +572,10 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, withRe if recording && recorded && db.Data.Serial > dmsDeployment.LastVersionID { return fmt.Errorf("inconsistent deployment state: resource state serial (%d) is ahead of deployment history version (%d)", db.Data.Serial, dmsDeployment.LastVersionID) } + _, shadow := db.Data.Features[FeatureDMSShadow] + if shadow && (recorded || recording) { + return errors.New("dms_shadow file state cannot be opened as DMS-authoritative state") + } switch { case recording && !recorded && len(db.Data.State) == 0: if db.Data.Features == nil { @@ -576,8 +591,6 @@ Run "databricks bundle destroy" first, then deploy again with deployment history return ErrUnsettingRecording } - db.recordsHistory = recorded - if recorded { // The service is the source of truth for a recorded deployment; the file is a tombstone // carrying only the marker, and applyDMSState loads the resources the service holds. @@ -612,44 +625,8 @@ Run "databricks bundle destroy" first, then deploy again with deployment history db.VersionID = dmsDeployment.LastVersionID } - if withWrite { - db.openedForWrite = true - - // A recorded deployment needs no WAL: every state write goes to the service as it happens, - // and Open recovers from there rather than replaying a log. Stamp the header fields the - // replay would have carried over, then skip the file. The serial is left alone: under - // recording it comes from the deployment's last_version_id. - if db.isDeploymentMetadataService() { - db.Data.Lineage = db.GetOrInitLineage() - db.Data.StateVersion = currentStateVersion - db.Data.CLIVersion = build.GetInfo().Version - return nil - } - - if err := os.MkdirAll(filepath.Dir(walPath), 0o755); err != nil { - return fmt.Errorf("failed to create state directory: %w", err) - } - walFile, err := os.OpenFile(walPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600) - if err != nil { - return fmt.Errorf("failed to open WAL file %s: %w", walPath, err) - } - db.walFile = walFile - walHead := Header{ - Lineage: db.GetOrInitLineage(), - Serial: db.Data.Serial + 1, - StateVersion: currentStateVersion, - CLIVersion: build.GetInfo().Version, - } - if err := appendJSONLine(db.walFile, walHead); err != nil { - // Remove the WAL created just above: O_EXCL means it is ours, and - // without a header it carries nothing to recover. Leaving it makes - // every later open fail, with recovery on the header parse and - // without it on the unexpected WAL file. - db.walFile.Close() - db.walFile = nil - os.Remove(walPath) - return fmt.Errorf("failed to write WAL header to %s: %w", walPath, err) - } + if mode.isWritable() { + return db.unlockedUpgradeToWrite(mode) } return nil @@ -666,6 +643,12 @@ func (db *DeploymentState) OpenWithData(path string, data Database) { } db.Path = path + db.mode = ReadingFile + db.operationBuffer = nil + db.versionCompleted = false + db.historyDrainTimeout = 0 + db.DeploymentID = "" + db.VersionID = 0 db.Data = data db.stateIDs = make(map[string]string) for key, entry := range data.State { @@ -714,6 +697,7 @@ func (db *DeploymentState) mergeWalIntoState(ctx context.Context) (bool, error) var ( newSerial int newCLIVersion string + newFeatures map[string]struct{} ) for scanner.Scan() { @@ -738,6 +722,16 @@ func (db *DeploymentState) mergeWalIntoState(ctx context.Context) (bool, error) if header.Serial > expectedSerial { return false, fmt.Errorf("WAL serial (%d) is ahead of expected (%d), state may be corrupted", header.Serial, expectedSerial) } + if err := assertNoUnsupportedFeatures(header.Features); err != nil { + return false, err + } + newFeatures = maps.Clone(db.Data.Features) + for feature := range header.Features { + if newFeatures == nil { + newFeatures = make(map[string]struct{}) + } + newFeatures[feature] = struct{}{} + } newSerial = header.Serial newCLIVersion = header.CLIVersion } else { @@ -774,11 +768,11 @@ func (db *DeploymentState) mergeWalIntoState(ctx context.Context) (bool, error) } } - hasEntries := lineNumber > 1 + hasEntries := lineNumber > 1 || lineNumber == 1 && !maps.Equal(db.Data.Features, newFeatures) - // Only advance the serial when the WAL carried entries, because the caller + // Only advance the serial when the WAL carried changes, because the caller // (replayWAL) persists the new state file only in that case. A header-only - // WAL is a deploy that started but committed nothing; advancing the serial + // WAL with unchanged features is a deploy that started but committed nothing; advancing the serial // for it leaves the in-memory serial ahead of the persisted one, so the // next deploy writes its WAL header at serial+2 and recovery rejects it as // "ahead of expected". See acceptance/bundle/deploy/wal/header-only-wal. @@ -790,6 +784,7 @@ func (db *DeploymentState) mergeWalIntoState(ctx context.Context) (bool, error) if hasEntries { db.Data.Serial = newSerial db.Data.CLIVersion = newCLIVersion + db.Data.Features = newFeatures } return hasEntries, nil @@ -815,7 +810,7 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported } db.walFile = nil err = db.replayWAL(ctx) - } else if db.openedForWrite && db.isDeploymentMetadataService() { + } else if db.mode == WritingDMS { // replayWAL is what normally persists the file. Without one, write the tombstone here so // the header (lineage, state version, CLI version) still lands on disk - even when no // operations were recorded, since the deployment and any version this run created exist @@ -823,10 +818,21 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported err = db.unlockedSave() } - // Wait until all operations are recorded in the service. - if db.isDeploymentMetadataService() && db.operationBuffer != nil { - if drainErr := db.operationBuffer.Drain(); drainErr != nil { - err = errors.Join(err, drainErr) + // File state is persisted before waiting on optional history recording. + if db.operationBuffer != nil { + drainCtx := ctx + cancel := func() {} + if db.mode.isFileWithHistory() { + drainCtx, cancel = context.WithTimeout(ctx, db.historyDrainTimeout) + } + drainErr := db.operationBuffer.Drain(drainCtx) + cancel() + if drainErr != nil { + if db.mode.isFileWithHistory() { + log.Warnf(ctx, "Failed to record deployment history: %s", drainErr) + } else { + err = errors.Join(err, drainErr) + } } } @@ -837,25 +843,40 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported return state, err } -// UpgradeToWrite transitions from read mode to write mode without re-reading state. -// State must already be open for read. This initializes the WAL for writing. -func (db *DeploymentState) UpgradeToWrite() error { +// UpgradeToWrite selects a write policy without re-reading resource state. +func (db *DeploymentState) UpgradeToWrite(mode Mode) error { db.mu.Lock() defer db.mu.Unlock() + return db.unlockedUpgradeToWrite(mode) +} +func (db *DeploymentState) unlockedUpgradeToWrite(mode Mode) error { if db.Path == "" { return errors.New("internal error: DeploymentState must be opened first") } - if db.openedForWrite { + if db.mode.isWritable() { return errors.New("internal error: DeploymentState is already open for write") } - db.openedForWrite = true + if !mode.isWritable() { + return fmt.Errorf("invalid deployment state write mode: %d", mode) + } + if err := assertNoUnsupportedFeatures(db.Data.Features); err != nil { + return err + } + _, recorded := db.Data.Features[FeatureDeploymentHistory] + _, shadow := db.Data.Features[FeatureDMSShadow] + if db.mode.usesDMSState() != mode.usesDMSState() || recorded != mode.usesDMSState() { + return errors.New("cannot change the deployment state source when upgrading to write") + } + if shadow && !mode.isFileWithHistory() { + return errors.New("dms_shadow state requires writing with deployment history enabled") + } - // As in Open: a recorded deployment writes no WAL, so stamp the header and skip the file. - if db.isDeploymentMetadataService() { + if mode == WritingDMS { db.Data.Lineage = db.GetOrInitLineage() db.Data.StateVersion = currentStateVersion db.Data.CLIVersion = build.GetInfo().Version + db.mode = mode return nil } @@ -867,15 +888,29 @@ func (db *DeploymentState) UpgradeToWrite() error { if err != nil { return fmt.Errorf("failed to open WAL file %s: %w", walPath, err) } - db.walFile = walFile - + features := maps.Clone(db.Data.Features) + if mode == WritingFileWithHistoryEnforced { + if features == nil { + features = make(map[string]struct{}) + } + features[FeatureDMSShadow] = struct{}{} + } walHead := Header{ Lineage: db.GetOrInitLineage(), Serial: db.Data.Serial + 1, StateVersion: currentStateVersion, CLIVersion: build.GetInfo().Version, + Features: features, + } + if err := appendJSONLine(walFile, walHead); err != nil { + // O_EXCL guarantees this is ours; an incomplete header cannot be recovered. + walFile.Close() + os.Remove(walPath) + return fmt.Errorf("failed to write WAL header to %s: %w", walPath, err) } - return appendJSONLine(db.walFile, walHead) + db.walFile = walFile + db.mode = mode + return nil } // Persist writes the current in-memory state to its file immediately, without a WAL cycle. @@ -903,14 +938,14 @@ func (db *DeploymentState) AssertOpenedForReadOrWrite() { func (db *DeploymentState) AssertOpenedForRead() { db.AssertOpenedForReadOrWrite() - if db.openedForWrite { + if db.mode.isWritable() { panic("internal error: DeploymentState must be opened in read mode") } } func (db *DeploymentState) AssertOpenedForWrite() { db.AssertOpenedForReadOrWrite() - if !db.openedForWrite { + if !db.mode.isWritable() { panic("internal error: DeploymentState must be opened in write mode") } } diff --git a/bundle/direct/dstate/state_test.go b/bundle/direct/dstate/state_test.go index 72ec263f7b7..6bf51304077 100644 --- a/bundle/direct/dstate/state_test.go +++ b/bundle/direct/dstate/state_test.go @@ -2,6 +2,7 @@ package dstate import ( "encoding/json" + "fmt" "os" "path/filepath" "testing" @@ -25,14 +26,14 @@ func TestOpenSaveFinalizeRoundTrip(t *testing.T) { path := filepath.Join(t.TempDir(), "state.json") var db DeploymentState - require.NoError(t, db.Open(t.Context(), path, WithRecovery(true), WithWrite(true), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, db.Open(t.Context(), path, WritingFile, WithRecovery(true), OpenDmsArgs{})) require.NoError(t, db.SaveState(t.Context(), "jobs.my_job", "123", map[string]string{"key": "val"}, nil)) mustFinalize(t, &db) // Re-open and verify persisted data. var db2 DeploymentState - require.NoError(t, db2.Open(t.Context(), path, WithRecovery(false), WithWrite(false), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, db2.Open(t.Context(), path, ReadingFile, WithRecovery(false), OpenDmsArgs{})) assert.Equal(t, 1, db2.Data.Serial) assert.Equal(t, "123", db2.GetResourceID("jobs.my_job")) mustFinalize(t, &db2) @@ -42,7 +43,7 @@ func TestFinalizeWithNoEntriesDoesNotWriteStateFile(t *testing.T) { path := filepath.Join(t.TempDir(), "state.json") var db DeploymentState - require.NoError(t, db.Open(t.Context(), path, WithRecovery(true), WithWrite(true), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, db.Open(t.Context(), path, WritingFile, WithRecovery(true), OpenDmsArgs{})) mustFinalize(t, &db) _, err := os.Stat(path) @@ -63,9 +64,12 @@ func TestDataForFileSerial(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { db := DeploymentState{ - Data: NewDatabase("test-lineage", 3), - recordsHistory: tt.recording, - VersionID: tt.versionID, + Data: NewDatabase("test-lineage", 3), + mode: ReadingFile, + VersionID: tt.versionID, + } + if tt.recording { + db.mode = ReadingDMS } db.Data.State["resources.jobs.foo"] = ResourceEntry{ID: "123", State: json.RawMessage(`{"name":"foo"}`)} @@ -179,7 +183,7 @@ func TestOpenUsesDMSVersionForSerial(t *testing.T) { ctx := cmdctx.SetWorkspaceClient(t.Context(), client) var db DeploymentState - err = db.Open(ctx, path, WithRecovery(false), WithWrite(false), WithDeploymentHistory(true), OpenDmsArgs{ + err = db.Open(ctx, path, ReadingDMS, WithRecovery(false), OpenDmsArgs{ DeploymentID: tt.deploymentID, LastVersionID: tt.lastVersionID, }) if tt.wantErr != "" { @@ -243,10 +247,10 @@ func TestPanicOnDoubleOpen(t *testing.T) { path := filepath.Join(t.TempDir(), "state.json") var db DeploymentState - require.NoError(t, db.Open(t.Context(), path, WithRecovery(true), WithWrite(true), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, db.Open(t.Context(), path, WritingFile, WithRecovery(true), OpenDmsArgs{})) assert.Panics(t, func() { - _ = db.Open(t.Context(), path, WithRecovery(true), WithWrite(true), WithDeploymentHistory(false), OpenDmsArgs{}) + _ = db.Open(t.Context(), path, WritingFile, WithRecovery(true), OpenDmsArgs{}) }) mustFinalize(t, &db) } @@ -264,12 +268,12 @@ func TestCLIVersionRecordsLastWriter(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte(seed), 0o600)) var db DeploymentState - require.NoError(t, db.Open(t.Context(), path, WithRecovery(true), WithWrite(true), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, db.Open(t.Context(), path, WritingFile, WithRecovery(true), OpenDmsArgs{})) require.NoError(t, db.SaveState(t.Context(), "resources.jobs.my_job", "123", map[string]string{"k": "v"}, nil)) mustFinalize(t, &db) var reopened DeploymentState - require.NoError(t, reopened.Open(t.Context(), path, WithRecovery(false), WithWrite(false), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, reopened.Open(t.Context(), path, ReadingFile, WithRecovery(false), OpenDmsArgs{})) assert.Equal(t, build.GetInfo().Version, reopened.Data.CLIVersion) assert.Equal(t, 2, reopened.Data.Serial) mustFinalize(t, &reopened) @@ -291,7 +295,7 @@ func TestHeaderOnlyWALDoesNotUpdateCLIVersion(t *testing.T) { require.NoError(t, os.WriteFile(walPath, append(headerLine, '\n'), 0o600)) var recovered DeploymentState - require.NoError(t, recovered.Open(t.Context(), path, WithRecovery(true), WithWrite(false), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, recovered.Open(t.Context(), path, ReadingFile, WithRecovery(true), OpenDmsArgs{})) assert.Equal(t, "0.1.2", recovered.Data.CLIVersion, "a header-only WAL wrote no state, so the version must not move") assert.Equal(t, 1, recovered.Data.Serial) mustFinalize(t, &recovered) @@ -303,12 +307,12 @@ func TestHeaderOnlyWALRecoveryDoesNotAdvanceSerial(t *testing.T) { // Commit serial 1 with one resource. var db DeploymentState - require.NoError(t, db.Open(t.Context(), path, WithRecovery(true), WithWrite(true), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, db.Open(t.Context(), path, WritingFile, WithRecovery(true), OpenDmsArgs{})) require.NoError(t, db.SaveState(t.Context(), "jobs.my_job", "123", map[string]string{}, nil)) mustFinalize(t, &db) var committed DeploymentState - require.NoError(t, committed.Open(t.Context(), path, WithRecovery(false), WithWrite(false), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, committed.Open(t.Context(), path, ReadingFile, WithRecovery(false), OpenDmsArgs{})) lineage := committed.Data.Lineage require.Equal(t, 1, committed.Data.Serial) mustFinalize(t, &committed) @@ -324,7 +328,7 @@ func TestHeaderOnlyWALRecoveryDoesNotAdvanceSerial(t *testing.T) { require.NoError(t, os.WriteFile(walPath, append(headerLine, '\n'), 0o600)) var recovered DeploymentState - require.NoError(t, recovered.Open(t.Context(), path, WithRecovery(true), WithWrite(false), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, recovered.Open(t.Context(), path, ReadingFile, WithRecovery(true), OpenDmsArgs{})) assert.Equal(t, 1, recovered.Data.Serial) assert.Equal(t, "123", recovered.GetResourceID("jobs.my_job")) assert.NoFileExists(t, walPath) @@ -353,17 +357,17 @@ func TestDeleteState(t *testing.T) { path := filepath.Join(t.TempDir(), "state.json") var db DeploymentState - require.NoError(t, db.Open(t.Context(), path, WithRecovery(true), WithWrite(true), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, db.Open(t.Context(), path, WritingFile, WithRecovery(true), OpenDmsArgs{})) require.NoError(t, db.SaveState(t.Context(), "jobs.my_job", "123", map[string]string{}, nil)) mustFinalize(t, &db) var db2 DeploymentState - require.NoError(t, db2.Open(t.Context(), path, WithRecovery(true), WithWrite(true), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, db2.Open(t.Context(), path, WritingFile, WithRecovery(true), OpenDmsArgs{})) require.NoError(t, db2.DeleteState(t.Context(), "jobs.my_job", false)) mustFinalize(t, &db2) var db3 DeploymentState - require.NoError(t, db3.Open(t.Context(), path, WithRecovery(false), WithWrite(false), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, db3.Open(t.Context(), path, ReadingFile, WithRecovery(false), OpenDmsArgs{})) assert.Equal(t, 2, db3.Data.Serial) assert.Empty(t, db3.GetResourceID("jobs.my_job")) mustFinalize(t, &db3) @@ -375,7 +379,7 @@ func TestGetOrInitLineageReadableBeforeWriteAndPersisted(t *testing.T) { // Fresh state opened read-only, as the deploy does before planning: no // lineage yet. var db DeploymentState - require.NoError(t, db.Open(t.Context(), path, WithRecovery(true), WithWrite(false), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, db.Open(t.Context(), path, ReadingFile, WithRecovery(true), OpenDmsArgs{})) require.Empty(t, db.Data.Lineage) // GetOrInitLineage initializes the lineage and makes it readable before any @@ -386,13 +390,13 @@ func TestGetOrInitLineageReadableBeforeWriteAndPersisted(t *testing.T) { // Upgrading to write reuses the same lineage (it goes into the WAL header), // and a write makes it durable. - require.NoError(t, db.UpgradeToWrite()) + require.NoError(t, db.UpgradeToWrite(WritingFile)) require.NoError(t, db.SaveState(t.Context(), "jobs.my_job", "123", map[string]string{}, nil)) mustFinalize(t, &db) // Re-open: the persisted lineage matches the one read before the write. var reopened DeploymentState - require.NoError(t, reopened.Open(t.Context(), path, WithRecovery(false), WithWrite(false), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, reopened.Open(t.Context(), path, ReadingFile, WithRecovery(false), OpenDmsArgs{})) assert.Equal(t, lineage, reopened.Data.Lineage) mustFinalize(t, &reopened) } @@ -406,13 +410,34 @@ func TestOpenFailureLeavesStateClosed(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte("{not json"), 0o600)) var db DeploymentState - require.Error(t, db.Open(t.Context(), path, WithRecovery(true), WithWrite(true), WithDeploymentHistory(false), OpenDmsArgs{})) + require.Error(t, db.Open(t.Context(), path, WritingFile, WithRecovery(true), OpenDmsArgs{})) assert.Empty(t, db.Path) // Once the state file is readable, the same receiver opens without panicking. seed := `{"state_version":2,"cli_version":"0.1.2","lineage":"test-lineage","serial":1,"state":{}}` require.NoError(t, os.WriteFile(path, []byte(seed), 0o600)) - require.NoError(t, db.Open(t.Context(), path, WithRecovery(true), WithWrite(true), WithDeploymentHistory(false), OpenDmsArgs{})) + require.NoError(t, db.Open(t.Context(), path, WritingFile, WithRecovery(true), OpenDmsArgs{})) assert.Equal(t, "test-lineage", db.Data.Lineage) mustFinalize(t, &db) } + +func TestFileWriteErrorsRemainFatalWithHistory(t *testing.T) { + for _, mode := range []Mode{WritingFile, WritingFileWithHistory, WritingFileWithHistoryEnforced} { + t.Run(fmt.Sprint(mode), func(t *testing.T) { + path := filepath.Join(t.TempDir(), "resources.json") + data := NewDatabase("lineage", 1) + data.State["resources.jobs.job"] = ResourceEntry{ID: "old-id"} + var db DeploymentState + db.OpenWithData(path, data) + db.mode = mode + wal, err := os.Create(path + walSuffix) + require.NoError(t, err) + require.NoError(t, wal.Close()) + db.walFile = wal + require.Error(t, db.SaveState(t.Context(), "resources.jobs.job", "new-id", struct{}{}, nil)) + assert.Equal(t, "old-id", db.GetResourceID("resources.jobs.job")) + require.Error(t, db.DeleteState(t.Context(), "resources.jobs.job", false)) + assert.Equal(t, "old-id", db.GetResourceID("resources.jobs.job")) + }) + } +} diff --git a/bundle/migrate/build_state_test.go b/bundle/migrate/build_state_test.go index a495436008c..f9c7cac9519 100644 --- a/bundle/migrate/build_state_test.go +++ b/bundle/migrate/build_state_test.go @@ -48,7 +48,7 @@ func runBuildStateFromTF( var db dstate.DeploymentState db.OpenWithData(statePath, dstate.NewDatabase("lineage", 1)) - require.NoError(t, db.UpgradeToWrite()) + require.NoError(t, db.UpgradeToWrite(dstate.WritingFile)) _, err = migrate.BuildStateFromTF(t.Context(), &root, adapters, &db, tfAttrs, tfIDs, "") require.NoError(t, err) diff --git a/bundle/phases/deploy.go b/bundle/phases/deploy.go index 13c5121f569..18e4ba4f4ef 100644 --- a/bundle/phases/deploy.go +++ b/bundle/phases/deploy.go @@ -15,6 +15,7 @@ import ( "github.com/databricks/cli/bundle/deploy/metadata" "github.com/databricks/cli/bundle/deploy/snapshot" "github.com/databricks/cli/bundle/deployplan" + "github.com/databricks/cli/bundle/direct/dstate" "github.com/databricks/cli/bundle/libraries" "github.com/databricks/cli/bundle/metrics" "github.com/databricks/cli/bundle/permissions" @@ -297,7 +298,11 @@ func Deploy(ctx context.Context, b *bundle.Bundle, outputHandler sync.OutputHand // Upgrade from read (opened by process.go, or by Migrate) to write mode. After approval // and the migration commit above, so a declined deploy never opens a WAL it must discard, // and the migration's tf+1 push lands before this advances the WAL header to tf+2. - if err := b.DeploymentBundle.StateDB.UpgradeToWrite(); err != nil { + writeMode := dstate.WritingFile + if b.DeploymentBundle.StateDB.IsDeploymentMetadataService() { + writeMode = dstate.WritingDMS + } + if err := b.DeploymentBundle.StateDB.UpgradeToWrite(writeMode); err != nil { logdiag.LogError(ctx, err) return } diff --git a/bundle/phases/destroy.go b/bundle/phases/destroy.go index 07bbbbca385..7b042a456b2 100644 --- a/bundle/phases/destroy.go +++ b/bundle/phases/destroy.go @@ -14,6 +14,7 @@ import ( "github.com/databricks/cli/bundle/deploy/files" "github.com/databricks/cli/bundle/deploy/lock" "github.com/databricks/cli/bundle/deployplan" + "github.com/databricks/cli/bundle/direct/dstate" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/diag" "github.com/databricks/cli/libs/dms" @@ -347,7 +348,11 @@ func Destroy(ctx context.Context, b *bundle.Bundle) { } // Upgrade from read (opened by process.go, or by Migrate) to write mode - if err := b.DeploymentBundle.StateDB.UpgradeToWrite(); err != nil { + writeMode := dstate.WritingFile + if b.DeploymentBundle.StateDB.IsDeploymentMetadataService() { + writeMode = dstate.WritingDMS + } + if err := b.DeploymentBundle.StateDB.UpgradeToWrite(writeMode); err != nil { logdiag.LogError(ctx, err) return } diff --git a/bundle/phases/dms.go b/bundle/phases/dms.go index a2399c9e3a1..e2d7ea958c1 100644 --- a/bundle/phases/dms.go +++ b/bundle/phases/dms.go @@ -179,7 +179,7 @@ func startVersion(ctx context.Context, b *bundle.Bundle, versionType dms.Version } log.Infof(ctx, "Created deployment version: deployment=%s version=%s", deploymentID, version.VersionId) - db.StateDB.InitializeOperationBuffer(ctx, deploymentID, versionID) + db.StateDB.InitializeOperationBuffer(ctx, deploymentID, versionID, len(staged)) return nil } diff --git a/bundle/statemgmt/direct_migration.go b/bundle/statemgmt/direct_migration.go index f3d632f138f..7a06d75aca9 100644 --- a/bundle/statemgmt/direct_migration.go +++ b/bundle/statemgmt/direct_migration.go @@ -196,7 +196,7 @@ func checkPlanOnTempState(ctx context.Context, b *bundle.Bundle, tempStatePath s // This plan is not created with the deployment history feature enabled, // so we can safely pass false for withDeploymentHistory. - if err := planBundle.StateDB.Open(planCtx, tempStatePath, false, false, dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{}); err != nil { + if err := planBundle.StateDB.Open(planCtx, tempStatePath, dstate.ReadingFile, dstate.WithRecovery(false), dstate.OpenDmsArgs{}); err != nil { return nil, fmt.Errorf("opening migrated state for plan check: %w", err) } @@ -287,7 +287,7 @@ func convertTFStateToDirect(ctx context.Context, b *bundle.Bundle, tfState *migr return tempStatePath, false, nil, err } - if err := stateDB.UpgradeToWrite(); err != nil { + if err := stateDB.UpgradeToWrite(dstate.WritingFile); err != nil { return tempStatePath, false, nil, fmt.Errorf("upgrading state for apply: %w", err) } diff --git a/cmd/bundle/deployment/migrate.go b/cmd/bundle/deployment/migrate.go index f329ec7962c..c9b82705d5e 100644 --- a/cmd/bundle/deployment/migrate.go +++ b/cmd/bundle/deployment/migrate.go @@ -149,7 +149,7 @@ New bundles use the direct engine automatically, no migration is needed.`) return err } - if err := stateDB.UpgradeToWrite(); err != nil { + if err := stateDB.UpgradeToWrite(dstate.WritingFile); err != nil { return fmt.Errorf("upgrading state for apply: %w", err) } diff --git a/cmd/bundle/utils/process.go b/cmd/bundle/utils/process.go index 21c3008f719..d542652b05e 100644 --- a/cmd/bundle/utils/process.go +++ b/cmd/bundle/utils/process.go @@ -358,7 +358,7 @@ func ProcessBundleRet(cmd *cobra.Command, opts ProcessOptions) (b *bundle.Bundle if needDirectState && !b.DeploymentBundle.StateDB.IsOpen() { _, localPath = b.StateFilenameDirect(ctx) if !stateDesc.IsDMS() { - if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WithRecovery(true), dstate.WithWrite(false), dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{}); err != nil { + if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.ReadingFile, dstate.WithRecovery(true), dstate.OpenDmsArgs{}); err != nil { logdiag.LogError(ctx, err) return b, stateDesc, root.ErrAlreadyPrinted } @@ -406,7 +406,7 @@ func ProcessBundleRet(cmd *cobra.Command, opts ProcessOptions) (b *bundle.Bundle if !cmdctx.HasWorkspaceClient(ctx) { ctx = cmdctx.SetWorkspaceClient(ctx, b.WorkspaceClient(ctx)) } - if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WithRecovery(false), dstate.WithWrite(false), dstate.WithDeploymentHistory(true), dstate.OpenDmsArgs{DeploymentID: dmsDeploymentID, LastVersionID: lastVersionID}); err != nil { + if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.ReadingDMS, dstate.WithRecovery(false), dstate.OpenDmsArgs{DeploymentID: dmsDeploymentID, LastVersionID: lastVersionID}); err != nil { logdiag.LogError(ctx, err) return b, stateDesc, root.ErrAlreadyPrinted } @@ -639,7 +639,7 @@ func parseLastVersionID(dmsDeployment *bundledeployments.Deployment) (int, error func OpenDirectStateForRead(ctx context.Context, b *bundle.Bundle, stateDesc *statemgmt.StateDesc) error { _, localPath := b.StateFilenameDirect(ctx) if !resolveDeploymentHistory(ctx, b, stateDesc) { - if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WithRecovery(true), dstate.WithWrite(false), dstate.WithDeploymentHistory(false), dstate.OpenDmsArgs{}); err != nil { + if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.ReadingFile, dstate.WithRecovery(true), dstate.OpenDmsArgs{}); err != nil { return err } return enforceDeploymentHistorySetting(ctx, b, stateDesc, false) @@ -657,7 +657,7 @@ func OpenDirectStateForRead(ctx context.Context, b *bundle.Bundle, stateDesc *st if !cmdctx.HasWorkspaceClient(ctx) { ctx = cmdctx.SetWorkspaceClient(ctx, b.WorkspaceClient(ctx)) } - if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WithRecovery(false), dstate.WithWrite(false), dstate.WithDeploymentHistory(true), dstate.OpenDmsArgs{DeploymentID: dmsDeploymentID, LastVersionID: lastVersionID}); err != nil { + if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.ReadingDMS, dstate.WithRecovery(false), dstate.OpenDmsArgs{DeploymentID: dmsDeploymentID, LastVersionID: lastVersionID}); err != nil { return err } return enforceDeploymentHistorySetting(ctx, b, stateDesc, false) diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index efecb3a330a..23f9f328ac4 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -3,6 +3,7 @@ package dms import ( "context" "encoding/json" + "errors" "fmt" "sync" @@ -14,6 +15,12 @@ import ( // plan reads. const bufferedOperations = 10 +// BufferOptions controls backpressure and whether recording stops after an error. +type BufferOptions struct { + Capacity int + BestEffort bool +} + // OperationBuffer records each state write with DMS for one deployment version, off the apply // path: writes are queued and sent on one background goroutine. It exists only while a bundle // records deployment history; callers hold a nil buffer otherwise and must not call it. @@ -23,9 +30,11 @@ type OperationBuffer struct { // The buffer. queue holds bundle state keys, and pending the newest update per key, so a // second write for a resource replaces the first. - queue chan string - done chan struct{} - stopQueue func() + queue chan string + done chan struct{} + stopQueue func() + cancel context.CancelFunc + bestEffort bool // sequenceIDs holds the token the last update for a resource returned. A missing key reads as // 0, which is what staging leaves. Unguarded: run is the only goroutine that writes, one update at a time. @@ -53,12 +62,19 @@ type OperationBuffer struct { // StartOperationBuffer opens the buffer for the version the caller just created. The version // must already exist: operations record under it, and nothing here creates it. -func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum int) *OperationBuffer { +func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum int, options BufferOptions) *OperationBuffer { + ctx, cancel := context.WithCancel(ctx) + capacity := options.Capacity + if capacity <= 0 { + capacity = bufferedOperations + } b := &OperationBuffer{ deploymentID: deploymentID, versionNum: versionNum, - queue: make(chan string, bufferedOperations), + queue: make(chan string, capacity), done: make(chan struct{}), + cancel: cancel, + bestEffort: options.BestEffort, pending: make(map[string]OperationUpdate), latestState: make(map[string]OperationUpdate), sequenceIDs: make(map[string]int64), @@ -69,7 +85,7 @@ func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum i } // RecordOperation records a state write. state is the serialized envelope, nil for a delete. -// An earlier failure does not stop it: keep recording, best effort. +// Required recording continues after an earlier failure; optional history stops on the first error. func (b *OperationBuffer) RecordOperation(ctx context.Context, resourceKey string, inProgress bool, resourceID string, state json.RawMessage) { update, err := NewStateUpdate(resourceID, state, inProgress) if err != nil { @@ -93,6 +109,10 @@ func (b *OperationBuffer) RecordFailure(resourceKey, resourceID string, state js // that latest state. latestState and pending are updated under the one lock. Panics after Drain. func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, keepExistingUpdateIfAny bool, cause error) { b.mu.Lock() + if b.bestEffort && b.err != nil { + b.mu.Unlock() + return + } if keepExistingUpdateIfAny { if latest, recorded := b.latestState[resourceKey]; recorded { update = latest.Merge(failureUpdate(cause)) @@ -112,7 +132,18 @@ func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, kee // // Only when the resource is not already queued, enqueue it. if !queued { - b.queue <- resourceKey + if b.bestEffort { + select { + case b.queue <- resourceKey: + default: + b.setErr(errors.New("deployment history buffer is full")) + } + } else { + select { + case b.queue <- resourceKey: + case <-b.done: + } + } } } @@ -128,8 +159,13 @@ func (b *OperationBuffer) take(resourceKey string) (OperationUpdate, bool) { func (b *OperationBuffer) run(ctx context.Context) { defer close(b.done) + defer b.cancel() for resourceKey := range b.queue { + if err := ctx.Err(); err != nil { + b.setErr(err) + return + } update, ok := b.take(resourceKey) if !ok { // Unreachable: a key is queued only when nothing was waiting for it. Guard so @@ -137,9 +173,12 @@ func (b *OperationBuffer) run(ctx context.Context) { continue } - // Keep going after a failure, so one bad write does not drop everything behind it. + // Required recording keeps going so one bad write does not drop everything behind it. if err := b.write(ctx, resourceKey, update); err != nil { b.setErr(fmt.Errorf("recording operation for %s: %w", resourceKey, err)) + if b.bestEffort { + return + } } } } @@ -161,12 +200,18 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio return nil } -// Drain waits for everything buffered to reach the service and returns the first write error, -// which fails the deploy: DMS is the source of truth for what exists. Safe to call twice. -func (b *OperationBuffer) Drain() error { +// Drain waits for buffered writes, canceling outstanding requests when ctx expires. +// It returns the first recording error and is safe to call twice. +func (b *OperationBuffer) Drain(ctx context.Context) error { b.stopQueue() - <-b.done - return b.Err() + select { + case <-b.done: + return b.Err() + case <-ctx.Done(): + b.setErr(ctx.Err()) + b.cancel() + return b.Err() + } } // setErr keeps the first error; one failure is enough to fail the deploy. @@ -176,6 +221,9 @@ func (b *OperationBuffer) setErr(err error) { if b.err == nil { b.err = err + if b.bestEffort { + b.cancel() + } } } diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index 49af18c3255..cc6589bea94 100644 --- a/libs/dms/operation_buffer_test.go +++ b/libs/dms/operation_buffer_test.go @@ -6,6 +6,7 @@ import ( "strings" "sync" "testing" + "time" "github.com/databricks/databricks-sdk-go/service/bundledeployments" "github.com/stretchr/testify/assert" @@ -21,6 +22,7 @@ import ( func queued() *OperationBuffer { return &OperationBuffer{ queue: make(chan string, bufferedOperations), + cancel: func() {}, pending: make(map[string]OperationUpdate), latestState: make(map[string]OperationUpdate), } @@ -76,8 +78,8 @@ func TestDrainIsIdempotent(t *testing.T) { s.stopQueue = sync.OnceFunc(func() { close(s.queue) }) go s.run(t.Context()) - require.NoError(t, s.Drain()) - require.NoError(t, s.Drain()) + require.NoError(t, s.Drain(t.Context())) + require.NoError(t, s.Drain(t.Context())) } func TestRecordFailureKeepsTheStateTheLastWriteRecorded(t *testing.T) { @@ -114,3 +116,33 @@ func TestRecordFailureUsesTheGivenStateWhenNoWriteRecorded(t *testing.T) { assert.JSONEq(t, `{"state":{"name":"v1"}}`, string(update.State)) assert.Equal(t, bundledeployments.OperationStatusOperationStatusFailed, update.Status) } + +func TestBestEffortBufferDoesNotBlockOnFullQueue(t *testing.T) { + b := queued() + b.bestEffort = true + for range bufferedOperations { + b.queue <- "occupied" + } + done := make(chan struct{}) + go func() { + b.RecordOperation(t.Context(), "resources.jobs.foo", false, "id-1", json.RawMessage(`{}`)) + close(done) + }() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("history recording blocked on a full queue") + } + assert.ErrorContains(t, b.Err(), "buffer is full") +} + +func TestBestEffortBufferStopsRecordingAfterError(t *testing.T) { + b := queued() + b.bestEffort = true + b.RecordOperation(t.Context(), "resources.jobs.foo", false, "id-1", json.RawMessage(strings.Repeat("x", maxStateSize+1))) + b.RecordOperation(t.Context(), "resources.jobs.bar", false, "id-2", json.RawMessage(`{}`)) + b.RecordFailure("resources.jobs.bar", "id-2", nil, errors.New("apply failed")) + assert.ErrorContains(t, b.Err(), "exceeds") + assert.Empty(t, b.queue) + assert.Empty(t, b.pending) +} From bbbf21007d5f6df1e8f68c4ab576a0f97a8711cc Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 10:57:18 +0000 Subject: [PATCH 02/52] Document deployment state modes and history error handling --- bundle/direct/dstate/mode.go | 18 ++++++++++++++++++ bundle/direct/dstate/state.go | 3 ++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/bundle/direct/dstate/mode.go b/bundle/direct/dstate/mode.go index 07fa8f61228..e4877184271 100644 --- a/bundle/direct/dstate/mode.go +++ b/bundle/direct/dstate/mode.go @@ -6,11 +6,29 @@ import "fmt" type Mode int const ( + // ReadingFile reads resources.json and optionally recovers a leftover WAL. + // State writes are disabled. ReadingFile Mode = iota + + // WritingFile records state changes in WAL. + // Finalize merges them into resources.json. DMS is unused. WritingFile + + // ReadingDMS reads resource state from DMS. + // State writes are disabled. ReadingDMS + + // WritingDMS records resource state in DMS without creating a WAL. + // DMS errors fail the deployment. WritingDMS + + // WritingFileWithHistory writes state through WAL to resources.json and records DMS history. + // DMS errors and drain timeouts let the deployment continue. WritingFileWithHistory + + // WritingFileWithHistoryEnforced writes state through WAL to resources.json and records DMS history. + // DMS errors and drain timeouts let the deployment continue. + // Persists dms_shadow to reject clients that cannot honor history recording. WritingFileWithHistoryEnforced ) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 8666f7c8cd0..f5f00eb09bf 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -94,7 +94,8 @@ type DeploymentState struct { mu sync.Mutex walFile *os.File - // mode survives Finalize so version completion retains its error policy. + // CompleteVersion runs after Finalize and still needs the mode. + // DMS state: errors fail deployment. File state with history: errors only warn. mode Mode // Maps resource key to ID. Unlike Data.State, this is up to date during writes (deploys). From cfa1615aa1c5780370451f79369afdfac7c56fc1 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 11:02:35 +0000 Subject: [PATCH 03/52] Use a fixed history drain timeout and clarify mode documentation --- bundle/direct/dstate/state.go | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index f5f00eb09bf..6541495b991 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -36,6 +36,7 @@ const ( initialBufferSize = 64 * 1024 maxWalEntrySize = 10 * 1024 * 1024 walSuffix = ".wal" + historyDrainTimeout = 5 * time.Second historyOperationTimeout = 5 * time.Second ) @@ -94,8 +95,8 @@ type DeploymentState struct { mu sync.Mutex walFile *os.File - // CompleteVersion runs after Finalize and still needs the mode. - // DMS state: errors fail deployment. File state with history: errors only warn. + // mode controls where resource state is read and written, whether DMS + // history is recorded, and whether DMS errors fail the deployment. mode Mode // Maps resource key to ID. Unlike Data.State, this is up to date during writes (deploys). @@ -109,8 +110,6 @@ type DeploymentState struct { // completion after an explicit one does nothing. versionCompleted bool - historyDrainTimeout time.Duration - // DeploymentID is the recorded deployment's id. Set at Open, or by the first deploy that creates // the deployment, which Open cannot know about. DeploymentID string @@ -188,7 +187,6 @@ func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploy if db.mode.isFileWithHistory() { options.BestEffort = true options.Capacity = operationCount - db.historyDrainTimeout = time.Duration(operationCount) * historyOperationTimeout } db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options) db.DeploymentID = deploymentID @@ -512,7 +510,6 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode M } db.operationBuffer = nil db.versionCompleted = false - db.historyDrainTimeout = 0 db.DeploymentID = dmsDeployment.DeploymentID db.VersionID = dmsDeployment.LastVersionID @@ -647,7 +644,6 @@ func (db *DeploymentState) OpenWithData(path string, data Database) { db.mode = ReadingFile db.operationBuffer = nil db.versionCompleted = false - db.historyDrainTimeout = 0 db.DeploymentID = "" db.VersionID = 0 db.Data = data @@ -824,7 +820,7 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported drainCtx := ctx cancel := func() {} if db.mode.isFileWithHistory() { - drainCtx, cancel = context.WithTimeout(ctx, db.historyDrainTimeout) + drainCtx, cancel = context.WithTimeout(ctx, historyDrainTimeout) } drainErr := db.operationBuffer.Drain(drainCtx) cancel() From 50ed19e3f883940bbcfbbdb8fc2347f01c4dfa49 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 11:16:09 +0000 Subject: [PATCH 04/52] Scale DMS history drain budget and cap requests at one minute --- bundle/direct/dstate/mode_test.go | 143 ++++++++++++++++++++++++------ bundle/direct/dstate/state.go | 28 +++--- libs/dms/operation_buffer.go | 6 ++ 3 files changed, 137 insertions(+), 40 deletions(-) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index 73d5e348a42..79cc6e00646 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -11,6 +11,7 @@ import ( "path/filepath" "sync" "testing" + "testing/synctest" "time" "github.com/databricks/cli/bundle/deployplan" @@ -32,6 +33,9 @@ func (m modeTransport) RoundTrip(r *http.Request) (*http.Response, error) { } else { m.handler.ServeHTTP(w, r) } + if err := r.Context().Err(); err != nil { + return nil, err + } response := w.Result() response.Request = r return response, nil @@ -43,6 +47,8 @@ func modeContext(t *testing.T, handler http.HandlerFunc) context.Context { client, err := databricks.NewWorkspaceClient(&databricks.Config{ Host: "https://state.test", Token: "test-token", AuthType: "pat", HTTPTransport: modeTransport{handler: handler}, + // Let our request deadline expire before the SDK's inactivity timeout. + HTTPTimeoutSeconds: 120, }) require.NoError(t, err) return cmdctx.SetWorkspaceClient(t.Context(), client) @@ -266,36 +272,117 @@ func TestHistoryErrorsFailOpen(t *testing.T) { func TestHistoryDrainTimeoutPreservesFileState(t *testing.T) { for _, mode := range []dstate.Mode{dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { + for _, tc := range []struct { + name string + operationCount int + callerTimeout time.Duration + wantDuration time.Duration + }{ + {name: "one operation", operationCount: 1, wantDuration: 3 * time.Second}, + {name: "all staged operations", operationCount: 4, wantDuration: 12 * time.Second}, + {name: "shorter caller deadline", operationCount: 4, callerTimeout: 20 * time.Millisecond, wantDuration: 20 * time.Millisecond}, + } { + t.Run(fmt.Sprintf("%d/%s", mode, tc.name), func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + started := make(chan struct{}) + finished := make(chan struct{}) + ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + close(started) + <-r.Context().Done() + close(finished) + }) + path := filepath.Join(t.TempDir(), "resources.json") + var db dstate.DeploymentState + require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) + db.InitializeOperationBuffer(ctx, "dep-1", 1, tc.operationCount) + // One buffered update must still get the budget for all staged operations. + require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) + <-started + drainCtx := ctx + if tc.callerTimeout > 0 { + var cancel context.CancelFunc + drainCtx, cancel = context.WithTimeout(ctx, tc.callerTimeout) + defer cancel() + } + start := time.Now() + _, err := db.Finalize(drainCtx) + require.NoError(t, err) + assert.Equal(t, tc.wantDuration, time.Since(start)) + assert.FileExists(t, path) + assert.NoFileExists(t, path+".wal") + <-finished + }) + }) + } + } +} + +func TestOperationTimeout(t *testing.T) { + for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { t.Run(fmt.Sprint(mode), func(t *testing.T) { - started := make(chan struct{}) - finished := make(chan struct{}) - ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { - _, _ = io.Copy(io.Discard, r.Body) - close(started) - <-r.Context().Done() - close(finished) + synctest.Test(t, func(t *testing.T) { + finished := make(chan struct{}) + ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + <-r.Context().Done() + close(finished) + }) + path := filepath.Join(t.TempDir(), "resources.json") + var db dstate.DeploymentState + require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) + db.InitializeOperationBuffer(ctx, "dep-1", 1, 100) + start := time.Now() + require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) + <-finished + synctest.Wait() + assert.Equal(t, time.Minute, time.Since(start)) + _, err := db.Finalize(ctx) + if mode == dstate.WritingDMS { + require.ErrorIs(t, err, context.DeadlineExceeded) + } else { + require.NoError(t, err) + var reopened dstate.DeploymentState + require.NoError(t, reopened.Open(ctx, path, dstate.ReadingFile, false, dstate.OpenDmsArgs{})) + assert.Equal(t, "id-1", reopened.GetResourceID("resources.jobs.job")) + _, err = reopened.Finalize(ctx) + require.NoError(t, err) + } + }) + }) + } +} + +func TestCompleteVersionTimeout(t *testing.T) { + for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { + t.Run(fmt.Sprint(mode), func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost { + _, _ = io.Copy(io.Discard, r.Body) + <-r.Context().Done() + } else { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"sequence_id":1}`)) + } + }) + var db dstate.DeploymentState + path := filepath.Join(t.TempDir(), "resources.json") + require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) + db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) + require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) + _, err := db.Finalize(ctx) + require.NoError(t, err) + start := time.Now() + completed, err := db.CompleteVersion(ctx, true) + assert.Equal(t, time.Minute, time.Since(start)) + assert.False(t, completed) + if mode == dstate.WritingDMS { + require.ErrorIs(t, err, context.DeadlineExceeded) + } else { + require.NoError(t, err) + } }) - path := filepath.Join(t.TempDir(), "resources.json") - var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) - db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) - require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) - select { - case <-started: - case <-time.After(5 * time.Second): - t.Fatal("DMS request did not start") - } - drainCtx, cancel := context.WithTimeout(ctx, 20*time.Millisecond) - defer cancel() - _, err := db.Finalize(drainCtx) - require.NoError(t, err) - assert.FileExists(t, path) - assert.NoFileExists(t, path+".wal") - select { - case <-finished: - case <-time.After(5 * time.Second): - t.Fatal("DMS request was not canceled") - } }) } } diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 6541495b991..4faf4f4c4c3 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -32,12 +32,11 @@ import ( const ( // currentStateVersion is the schema version written by this CLI and the version // legacy states are migrated up to on load. - currentStateVersion = 3 - initialBufferSize = 64 * 1024 - maxWalEntrySize = 10 * 1024 * 1024 - walSuffix = ".wal" - historyDrainTimeout = 5 * time.Second - historyOperationTimeout = 5 * time.Second + currentStateVersion = 3 + initialBufferSize = 64 * 1024 + maxWalEntrySize = 10 * 1024 * 1024 + walSuffix = ".wal" + historyDrainTimeoutPerOperation = 3 * time.Second ) // FeatureDeploymentHistory marks a state whose resources are also recorded with the @@ -106,6 +105,9 @@ type DeploymentState struct { // history, in which case InitializeOperationBuffer installs it once the version exists. operationBuffer *dms.OperationBuffer + // Includes all staged operations, even those already recorded or coalesced. + operationCount int + // versionCompleted makes CompleteVersion a no-op after the first call, so a deferred safety-net // completion after an explicit one does nothing. versionCompleted bool @@ -189,6 +191,7 @@ func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploy options.Capacity = operationCount } db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options) + db.operationCount = operationCount db.DeploymentID = deploymentID db.VersionID = versionID } @@ -231,11 +234,8 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b bestEffort := db.mode.isFileWithHistory() db.mu.Unlock() - if bestEffort { - var cancel context.CancelFunc - ctx, cancel = context.WithTimeout(ctx, historyOperationTimeout) - defer cancel() - } + ctx, cancel := context.WithTimeout(ctx, dms.OperationTimeout) + defer cancel() // A recording failure fails the version even when the caller counted the deploy a success: the // service does not then hold everything the WAL does. Finalize already drained and surfaced it; @@ -497,6 +497,7 @@ func (db *DeploymentState) reset() { db.Path = "" db.Data = Database{} db.stateIDs = nil + db.operationCount = 0 } func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode Mode, withRecovery WithRecovery, dmsDeployment OpenDmsArgs) error { @@ -510,6 +511,7 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode M } db.operationBuffer = nil db.versionCompleted = false + db.operationCount = 0 db.DeploymentID = dmsDeployment.DeploymentID db.VersionID = dmsDeployment.LastVersionID @@ -644,6 +646,7 @@ func (db *DeploymentState) OpenWithData(path string, data Database) { db.mode = ReadingFile db.operationBuffer = nil db.versionCompleted = false + db.operationCount = 0 db.DeploymentID = "" db.VersionID = 0 db.Data = data @@ -820,7 +823,8 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported drainCtx := ctx cancel := func() {} if db.mode.isFileWithHistory() { - drainCtx, cancel = context.WithTimeout(ctx, historyDrainTimeout) + drainTimeout := time.Duration(db.operationCount) * historyDrainTimeoutPerOperation + drainCtx, cancel = context.WithTimeout(ctx, drainTimeout) } drainErr := db.operationBuffer.Drain(drainCtx) cancel() diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 23f9f328ac4..4ca916d9aa7 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "sync" + "time" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/databricks-sdk-go/service/bundledeployments" @@ -15,6 +16,9 @@ import ( // plan reads. const bufferedOperations = 10 +// OperationTimeout caps a single DMS history request, including its retries. +const OperationTimeout = time.Minute + // BufferOptions controls backpressure and whether recording stops after an error. type BufferOptions struct { Capacity int @@ -185,6 +189,8 @@ func (b *OperationBuffer) run(ctx context.Context) { // write sends one update, at the sequence id the resource is at. func (b *OperationBuffer) write(ctx context.Context, key string, update OperationUpdate) error { + ctx, cancel := context.WithTimeout(ctx, OperationTimeout) + defer cancel() w := cmdctx.WorkspaceClient(ctx) result, err := w.BundleDeployments.UpdateOperation(ctx, bundledeployments.UpdateOperationRequest{ Name: OperationName(b.deploymentID, b.versionNum, key), From aa5e06a07a57c215a9a732c53df313a786cda426 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 11:29:20 +0000 Subject: [PATCH 05/52] Set history drain budget to five seconds per operation --- bundle/direct/dstate/mode_test.go | 4 ++-- bundle/direct/dstate/state.go | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index 79cc6e00646..06ca14f8f06 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -278,8 +278,8 @@ func TestHistoryDrainTimeoutPreservesFileState(t *testing.T) { callerTimeout time.Duration wantDuration time.Duration }{ - {name: "one operation", operationCount: 1, wantDuration: 3 * time.Second}, - {name: "all staged operations", operationCount: 4, wantDuration: 12 * time.Second}, + {name: "one operation", operationCount: 1, wantDuration: 5 * time.Second}, + {name: "all staged operations", operationCount: 4, wantDuration: 20 * time.Second}, {name: "shorter caller deadline", operationCount: 4, callerTimeout: 20 * time.Millisecond, wantDuration: 20 * time.Millisecond}, } { t.Run(fmt.Sprintf("%d/%s", mode, tc.name), func(t *testing.T) { diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 4faf4f4c4c3..0bf43a9586c 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -36,7 +36,7 @@ const ( initialBufferSize = 64 * 1024 maxWalEntrySize = 10 * 1024 * 1024 walSuffix = ".wal" - historyDrainTimeoutPerOperation = 3 * time.Second + historyDrainTimeoutPerOperation = 5 * time.Second ) // FeatureDeploymentHistory marks a state whose resources are also recorded with the From f734f165adc848faf54d402602de2b2d4ae4e461 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 11:48:12 +0000 Subject: [PATCH 06/52] Rename best-effort history feature and store computed drain timeout --- bundle/direct/dstate/mode.go | 2 +- bundle/direct/dstate/mode_test.go | 28 +++++++++++------------ bundle/direct/dstate/state.go | 37 +++++++++++++++---------------- 3 files changed, 33 insertions(+), 34 deletions(-) diff --git a/bundle/direct/dstate/mode.go b/bundle/direct/dstate/mode.go index e4877184271..910370efc9e 100644 --- a/bundle/direct/dstate/mode.go +++ b/bundle/direct/dstate/mode.go @@ -28,7 +28,7 @@ const ( // WritingFileWithHistoryEnforced writes state through WAL to resources.json and records DMS history. // DMS errors and drain timeouts let the deployment continue. - // Persists dms_shadow to reject clients that cannot honor history recording. + // Persists deployment_history_best_effort to reject clients that cannot honor history recording. WritingFileWithHistoryEnforced ) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index 06ca14f8f06..b11f4b32638 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -137,7 +137,7 @@ func TestFileModesPersistStateAndHistory(t *testing.T) { assert.Equal(t, 1, reopened.GetSerial()) assert.Len(t, reopened.Data.State, 1) assert.Equal(t, dependsOn, reopened.Data.State["resources.jobs.kept"].DependsOn) - _, enforced := reopened.Data.Features[dstate.FeatureDMSShadow] + _, enforced := reopened.Data.Features[dstate.FeatureDeploymentHistoryBestEffort] assert.Equal(t, mode == dstate.WritingFileWithHistoryEnforced, enforced) _, err = reopened.Finalize(t.Context()) require.NoError(t, err) @@ -190,11 +190,11 @@ func TestDMSReadsAndWritesWithoutWAL(t *testing.T) { func TestUpgradeToWriteValidatesMode(t *testing.T) { for _, tc := range []struct { - name string - read dstate.Mode - write dstate.Mode - shadow bool - wantErr string + name string + read dstate.Mode + write dstate.Mode + bestEffort bool + wantErr string }{ {"file", dstate.ReadingFile, dstate.WritingFile, false, ""}, {"history", dstate.ReadingFile, dstate.WritingFileWithHistory, false, ""}, @@ -206,13 +206,13 @@ func TestUpgradeToWriteValidatesMode(t *testing.T) { {"reading mode", dstate.ReadingFile, dstate.ReadingFile, false, "invalid"}, {"unknown mode", dstate.ReadingFile, dstate.Mode(99), false, "invalid"}, {"enforcement retained", dstate.ReadingFile, dstate.WritingFileWithHistory, true, ""}, - {"enforcement bypass", dstate.ReadingFile, dstate.WritingFile, true, "dms_shadow"}, + {"enforcement bypass", dstate.ReadingFile, dstate.WritingFile, true, "deployment_history_best_effort"}, } { t.Run(tc.name, func(t *testing.T) { ctx := modeContext(t, nil) path := filepath.Join(t.TempDir(), "resources.json") - if tc.shadow { - require.NoError(t, os.WriteFile(path, []byte(`{"state_version":3,"features":{"dms_shadow":{}},"state":{}}`), 0o600)) + if tc.bestEffort { + require.NoError(t, os.WriteFile(path, []byte(`{"state_version":3,"features":{"deployment_history_best_effort":{}},"state":{}}`), 0o600)) } var db dstate.DeploymentState require.NoError(t, db.Open(ctx, path, tc.read, false, dstate.OpenDmsArgs{})) @@ -406,11 +406,11 @@ func TestEnforcedHistorySurvivesWALRecovery(t *testing.T) { require.NoError(t, os.WriteFile(recoveredPath+".wal", wal, 0o600)) var recovered dstate.DeploymentState require.NoError(t, recovered.Open(t.Context(), recoveredPath, dstate.ReadingFile, true, dstate.OpenDmsArgs{})) - assert.Contains(t, recovered.StateFeatures(), dstate.FeatureDMSShadow) + assert.Contains(t, recovered.StateFeatures(), dstate.FeatureDeploymentHistoryBestEffort) assert.Equal(t, 1, recovered.GetSerial()) assert.FileExists(t, recoveredPath) assert.NoFileExists(t, recoveredPath+".wal") - require.ErrorContains(t, recovered.UpgradeToWrite(dstate.WritingFile), "dms_shadow") + require.ErrorContains(t, recovered.UpgradeToWrite(dstate.WritingFile), "deployment_history_best_effort") if withResource { assert.Equal(t, "id-1", recovered.GetResourceID("resources.jobs.job")) } @@ -429,10 +429,10 @@ func TestOpenRejectsInvalidModesAndFeatures(t *testing.T) { }{ {"unknown mode", dstate.Mode(99), `{}`, "invalid deployment state mode"}, {"unknown feature", dstate.ReadingFile, `{"future":{}}`, "requires features"}, - {"conflicting features", dstate.ReadingFile, `{"deployment_history":{},"dms_shadow":{}}`, "dms_shadow"}, - {"DMS reading shadow", dstate.ReadingDMS, `{"dms_shadow":{}}`, "dms_shadow"}, + {"conflicting features", dstate.ReadingFile, `{"deployment_history":{},"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, + {"DMS reading best effort history", dstate.ReadingDMS, `{"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, {"file writing DMS", dstate.WritingFile, `{"deployment_history":{}}`, "not supported"}, - {"file writing shadow", dstate.WritingFile, `{"dms_shadow":{}}`, "dms_shadow"}, + {"file writing best effort history", dstate.WritingFile, `{"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, } { t.Run(tc.name, func(t *testing.T) { path := filepath.Join(t.TempDir(), "resources.json") diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 0bf43a9586c..90b6192220b 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -49,8 +49,9 @@ const ( // deploying such a state without recording is refused. const FeatureDeploymentHistory = "deployment_history" -// FeatureDMSShadow requires writers to record history while keeping file state authoritative. -const FeatureDMSShadow = "dms_shadow" +// FeatureDeploymentHistoryBestEffort requires writers to attempt DMS history recording. +// Resource state stays in resources.json; recording errors only warn. +const FeatureDeploymentHistoryBestEffort = "deployment_history_best_effort" // featuresDocURL explains the state-features mechanism in the error shown when this CLI refuses a // state that depends on a feature it does not recognize. @@ -60,8 +61,8 @@ const featuresDocURL = "https://docs.databricks.com/aws/en/dev-tools/bundles/sta // any feature not listed here is refused (see checkStateFeatures), so a newer CLI's feature is not // silently clobbered by this one. var recognizedFeatures = map[string]struct{}{ - FeatureDeploymentHistory: {}, - FeatureDMSShadow: {}, + FeatureDeploymentHistory: {}, + FeatureDeploymentHistoryBestEffort: {}, } // assertNoUnsupportedFeatures refuses a state that depends on a feature this CLI does not recognize. @@ -105,8 +106,7 @@ type DeploymentState struct { // history, in which case InitializeOperationBuffer installs it once the version exists. operationBuffer *dms.OperationBuffer - // Includes all staged operations, even those already recorded or coalesced. - operationCount int + historyDrainTimeout time.Duration // versionCompleted makes CompleteVersion a no-op after the first call, so a deferred safety-net // completion after an explicit one does nothing. @@ -191,7 +191,7 @@ func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploy options.Capacity = operationCount } db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options) - db.operationCount = operationCount + db.historyDrainTimeout = time.Duration(operationCount) * historyDrainTimeoutPerOperation db.DeploymentID = deploymentID db.VersionID = versionID } @@ -497,7 +497,7 @@ func (db *DeploymentState) reset() { db.Path = "" db.Data = Database{} db.stateIDs = nil - db.operationCount = 0 + db.historyDrainTimeout = 0 } func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode Mode, withRecovery WithRecovery, dmsDeployment OpenDmsArgs) error { @@ -511,7 +511,7 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode M } db.operationBuffer = nil db.versionCompleted = false - db.operationCount = 0 + db.historyDrainTimeout = 0 db.DeploymentID = dmsDeployment.DeploymentID db.VersionID = dmsDeployment.LastVersionID @@ -572,9 +572,9 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode M if recording && recorded && db.Data.Serial > dmsDeployment.LastVersionID { return fmt.Errorf("inconsistent deployment state: resource state serial (%d) is ahead of deployment history version (%d)", db.Data.Serial, dmsDeployment.LastVersionID) } - _, shadow := db.Data.Features[FeatureDMSShadow] - if shadow && (recorded || recording) { - return errors.New("dms_shadow file state cannot be opened as DMS-authoritative state") + _, bestEffort := db.Data.Features[FeatureDeploymentHistoryBestEffort] + if bestEffort && (recorded || recording) { + return errors.New("deployment_history_best_effort file state cannot be opened as DMS-authoritative state") } switch { case recording && !recorded && len(db.Data.State) == 0: @@ -646,7 +646,7 @@ func (db *DeploymentState) OpenWithData(path string, data Database) { db.mode = ReadingFile db.operationBuffer = nil db.versionCompleted = false - db.operationCount = 0 + db.historyDrainTimeout = 0 db.DeploymentID = "" db.VersionID = 0 db.Data = data @@ -823,8 +823,7 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported drainCtx := ctx cancel := func() {} if db.mode.isFileWithHistory() { - drainTimeout := time.Duration(db.operationCount) * historyDrainTimeoutPerOperation - drainCtx, cancel = context.WithTimeout(ctx, drainTimeout) + drainCtx, cancel = context.WithTimeout(ctx, db.historyDrainTimeout) } drainErr := db.operationBuffer.Drain(drainCtx) cancel() @@ -865,12 +864,12 @@ func (db *DeploymentState) unlockedUpgradeToWrite(mode Mode) error { return err } _, recorded := db.Data.Features[FeatureDeploymentHistory] - _, shadow := db.Data.Features[FeatureDMSShadow] + _, bestEffort := db.Data.Features[FeatureDeploymentHistoryBestEffort] if db.mode.usesDMSState() != mode.usesDMSState() || recorded != mode.usesDMSState() { return errors.New("cannot change the deployment state source when upgrading to write") } - if shadow && !mode.isFileWithHistory() { - return errors.New("dms_shadow state requires writing with deployment history enabled") + if bestEffort && !mode.isFileWithHistory() { + return errors.New("deployment_history_best_effort state requires writing with deployment history enabled") } if mode == WritingDMS { @@ -894,7 +893,7 @@ func (db *DeploymentState) unlockedUpgradeToWrite(mode Mode) error { if features == nil { features = make(map[string]struct{}) } - features[FeatureDMSShadow] = struct{}{} + features[FeatureDeploymentHistoryBestEffort] = struct{}{} } walHead := Header{ Lineage: db.GetOrInitLineage(), From b9f937ef9abea34678179d5b46528a394c40181a Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 11:57:44 +0000 Subject: [PATCH 07/52] Use a fixed DMS buffer and wait for queue space --- bundle/direct/dstate/mode_test.go | 30 ++++++++++++++ bundle/direct/dstate/state.go | 6 +-- libs/dms/operation_buffer.go | 24 +++-------- libs/dms/operation_buffer_test.go | 66 ++++++++++++++++++++++--------- 4 files changed, 84 insertions(+), 42 deletions(-) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index b11f4b32638..647d9f3d3b0 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -9,6 +9,7 @@ import ( "net/http/httptest" "os" "path/filepath" + "strconv" "sync" "testing" "testing/synctest" @@ -506,3 +507,32 @@ func TestRecoveryRejectsUnknownWALFeatures(t *testing.T) { assert.NoFileExists(t, path) assert.FileExists(t, path+".wal") } + +func TestFullHistoryBufferUnblocksAfterOperationTimeout(t *testing.T) { + for _, mode := range []dstate.Mode{dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { + t.Run(fmt.Sprint(mode), func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(io.Discard, r.Body) + <-r.Context().Done() + }) + path := filepath.Join(t.TempDir(), "resources.json") + var db dstate.DeploymentState + require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) + const operationCount = 32 + db.InitializeOperationBuffer(ctx, "dep-1", 1, operationCount) + start := time.Now() + for i := range operationCount { + key := fmt.Sprintf("resources.jobs.job_%d", i) + require.NoError(t, db.SaveState(ctx, key, strconv.Itoa(i), struct{}{}, nil)) + } + assert.Equal(t, time.Minute, time.Since(start)) + exported, err := db.Finalize(ctx) + require.NoError(t, err) + assert.Len(t, exported, operationCount) + assert.FileExists(t, path) + assert.NoFileExists(t, path+".wal") + }) + }) + } +} diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 90b6192220b..bfd8e7125b3 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -185,11 +185,7 @@ func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploy db.AssertOpenedForWrite() db.mu.Lock() defer db.mu.Unlock() - options := dms.BufferOptions{} - if db.mode.isFileWithHistory() { - options.BestEffort = true - options.Capacity = operationCount - } + options := dms.BufferOptions{BestEffort: db.mode.isFileWithHistory()} db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options) db.historyDrainTimeout = time.Duration(operationCount) * historyDrainTimeoutPerOperation db.DeploymentID = deploymentID diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 4ca916d9aa7..d3d7e7de81b 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -3,7 +3,6 @@ package dms import ( "context" "encoding/json" - "errors" "fmt" "sync" "time" @@ -19,9 +18,8 @@ const bufferedOperations = 10 // OperationTimeout caps a single DMS history request, including its retries. const OperationTimeout = time.Minute -// BufferOptions controls backpressure and whether recording stops after an error. +// BufferOptions controls whether recording continues after an error. type BufferOptions struct { - Capacity int BestEffort bool } @@ -68,14 +66,10 @@ type OperationBuffer struct { // must already exist: operations record under it, and nothing here creates it. func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum int, options BufferOptions) *OperationBuffer { ctx, cancel := context.WithCancel(ctx) - capacity := options.Capacity - if capacity <= 0 { - capacity = bufferedOperations - } b := &OperationBuffer{ deploymentID: deploymentID, versionNum: versionNum, - queue: make(chan string, capacity), + queue: make(chan string, bufferedOperations), done: make(chan struct{}), cancel: cancel, bestEffort: options.BestEffort, @@ -136,17 +130,9 @@ func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, kee // // Only when the resource is not already queued, enqueue it. if !queued { - if b.bestEffort { - select { - case b.queue <- resourceKey: - default: - b.setErr(errors.New("deployment history buffer is full")) - } - } else { - select { - case b.queue <- resourceKey: - case <-b.done: - } + select { + case b.queue <- resourceKey: + case <-b.done: } } } diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index cc6589bea94..bb7259602df 100644 --- a/libs/dms/operation_buffer_test.go +++ b/libs/dms/operation_buffer_test.go @@ -6,7 +6,7 @@ import ( "strings" "sync" "testing" - "time" + "testing/synctest" "github.com/databricks/databricks-sdk-go/service/bundledeployments" "github.com/stretchr/testify/assert" @@ -117,23 +117,53 @@ func TestRecordFailureUsesTheGivenStateWhenNoWriteRecorded(t *testing.T) { assert.Equal(t, bundledeployments.OperationStatusOperationStatusFailed, update.Status) } -func TestBestEffortBufferDoesNotBlockOnFullQueue(t *testing.T) { - b := queued() - b.bestEffort = true - for range bufferedOperations { - b.queue <- "occupied" - } - done := make(chan struct{}) - go func() { - b.RecordOperation(t.Context(), "resources.jobs.foo", false, "id-1", json.RawMessage(`{}`)) - close(done) - }() - select { - case <-done: - case <-time.After(time.Second): - t.Fatal("history recording blocked on a full queue") - } - assert.ErrorContains(t, b.Err(), "buffer is full") +func TestBestEffortBufferWaitsForQueueSpace(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + b := queued() + b.bestEffort = true + for range bufferedOperations { + b.queue <- "occupied" + } + done := make(chan struct{}) + go func() { + b.RecordOperation(t.Context(), "resources.jobs.foo", false, "id-1", json.RawMessage(`{}`)) + close(done) + }() + synctest.Wait() + select { + case <-done: + t.Fatal("recording should wait for queue space") + default: + } + <-b.queue + <-done + assert.NoError(t, b.Err()) + assert.Len(t, b.queue, bufferedOperations) + _, ok := b.take("resources.jobs.foo") + assert.True(t, ok) + }) +} + +func TestBestEffortBufferUnblocksWhenWriterStops(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + b := queued() + b.bestEffort = true + b.done = make(chan struct{}) + for range bufferedOperations { + b.queue <- "occupied" + } + done := make(chan struct{}) + go func() { + b.RecordOperation(t.Context(), "resources.jobs.foo", false, "id-1", json.RawMessage(`{}`)) + close(done) + }() + synctest.Wait() + cause := errors.New("recording failed") + b.setErr(cause) + close(b.done) + <-done + assert.ErrorIs(t, b.Err(), cause) + }) } func TestBestEffortBufferStopsRecordingAfterError(t *testing.T) { From 4cd4e0847f4fdd3cdffe0cbb5e9891e687097846 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 12:14:35 +0000 Subject: [PATCH 08/52] Clarify the file-mode guard in version completion --- bundle/direct/dstate/state.go | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index bfd8e7125b3..f1a1efe24a9 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -221,7 +221,12 @@ func (db *DeploymentState) RecordingError() error { func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (bool, error) { db.mu.Lock() buf := db.operationBuffer - if buf == nil || db.versionCompleted { + // WritingFile has no DMS buffer, so version completion is skipped. + if buf == nil { + db.mu.Unlock() + return false, nil + } + if db.versionCompleted { db.mu.Unlock() return false, nil } From 34a407075f2998250e3a76788abd8da141f1c89f Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 12:22:51 +0000 Subject: [PATCH 09/52] Use direct enum comparisons for deployment state modes --- bundle/direct/dstate/mode.go | 12 ---------- bundle/direct/dstate/state.go | 41 +++++++++++++++++++---------------- 2 files changed, 22 insertions(+), 31 deletions(-) diff --git a/bundle/direct/dstate/mode.go b/bundle/direct/dstate/mode.go index 910370efc9e..5280af1db39 100644 --- a/bundle/direct/dstate/mode.go +++ b/bundle/direct/dstate/mode.go @@ -41,18 +41,6 @@ func (m Mode) isWritable() bool { } } -func (m Mode) usesDMSState() bool { - return m == ReadingDMS || m == WritingDMS -} - -func (m Mode) isFileWithHistory() bool { - return m == WritingFileWithHistory || m == WritingFileWithHistoryEnforced -} - -func (m Mode) recordsHistory() bool { - return m.usesDMSState() || m.isFileWithHistory() -} - func (m Mode) validate() error { switch m { case ReadingFile, WritingFile, ReadingDMS, WritingDMS, WritingFileWithHistory, WritingFileWithHistoryEnforced: diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index f1a1efe24a9..e83eef71414 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -185,7 +185,9 @@ func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploy db.AssertOpenedForWrite() db.mu.Lock() defer db.mu.Unlock() - options := dms.BufferOptions{BestEffort: db.mode.isFileWithHistory()} + options := dms.BufferOptions{ + BestEffort: db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced, + } db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options) db.historyDrainTimeout = time.Duration(operationCount) * historyDrainTimeoutPerOperation db.DeploymentID = deploymentID @@ -232,7 +234,7 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b } db.versionCompleted = true deploymentID, versionID := db.DeploymentID, db.VersionID - bestEffort := db.mode.isFileWithHistory() + bestEffort := db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced db.mu.Unlock() ctx, cancel := context.WithTimeout(ctx, dms.OperationTimeout) @@ -324,7 +326,7 @@ func (db *DeploymentState) SaveState(ctx context.Context, key, newID string, sta DependsOn: dependsOn, } - if db.mode.usesDMSState() { + if db.mode == WritingDMS { db.Data.State[key] = entry } else if err := appendJSONLine(db.walFile, WALEntry{Key: key, Value: &entry}); err != nil { return err @@ -356,7 +358,7 @@ func (db *DeploymentState) DeleteState(ctx context.Context, key string, inProgre // Read before the delete below: DMS needs the id to say which resource went away. deletedID := db.stateIDs[key] - if db.mode.usesDMSState() { + if db.mode == WritingDMS { delete(db.Data.State, key) } else if err := appendJSONLine(db.walFile, WALEntry{Key: key}); err != nil { return err @@ -419,31 +421,31 @@ func (db *DeploymentState) StateFeatures() map[string]struct{} { func (db *DeploymentState) GetSerial() int { db.mu.Lock() defer db.mu.Unlock() - if db.isDeploymentMetadataService() { + if db.mode == ReadingDMS || db.mode == WritingDMS { return db.VersionID } return db.Data.Serial } -// isDeploymentMetadataService is IsDeploymentMetadataService for callers already holding db.mu. -func (db *DeploymentState) isDeploymentMetadataService() bool { - return db.mode.usesDMSState() -} - // IsDeploymentMetadataService reports whether this deployment's resource state lives in the // deployment metadata service rather than the state file. The state's features are the source of // truth; valid only after Open. func (db *DeploymentState) IsDeploymentMetadataService() bool { db.mu.Lock() defer db.mu.Unlock() - return db.isDeploymentMetadataService() + return db.mode == ReadingDMS || db.mode == WritingDMS } // RecordsHistory reports whether this mode records resource operations with DMS. func (db *DeploymentState) RecordsHistory() bool { db.mu.Lock() defer db.mu.Unlock() - return db.mode.recordsHistory() + switch db.mode { + case ReadingDMS, WritingDMS, WritingFileWithHistory, WritingFileWithHistoryEnforced: + return true + default: + return false + } } // GetOrInitLineage returns the deployment lineage, generating and storing a new @@ -507,7 +509,7 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode M } db.Path = path db.mode = ReadingFile - if mode.usesDMSState() { + if mode == ReadingDMS || mode == WritingDMS { db.mode = ReadingDMS } db.operationBuffer = nil @@ -540,7 +542,7 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode M db.stateIDs[key] = entry.ID } - recording := mode.usesDMSState() + recording := mode == ReadingDMS || mode == WritingDMS walPath := db.Path + walSuffix _, err = os.Stat(walPath) @@ -821,15 +823,16 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported // File state is persisted before waiting on optional history recording. if db.operationBuffer != nil { + bestEffort := db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced drainCtx := ctx cancel := func() {} - if db.mode.isFileWithHistory() { + if bestEffort { drainCtx, cancel = context.WithTimeout(ctx, db.historyDrainTimeout) } drainErr := db.operationBuffer.Drain(drainCtx) cancel() if drainErr != nil { - if db.mode.isFileWithHistory() { + if bestEffort { log.Warnf(ctx, "Failed to record deployment history: %s", drainErr) } else { err = errors.Join(err, drainErr) @@ -866,10 +869,10 @@ func (db *DeploymentState) unlockedUpgradeToWrite(mode Mode) error { } _, recorded := db.Data.Features[FeatureDeploymentHistory] _, bestEffort := db.Data.Features[FeatureDeploymentHistoryBestEffort] - if db.mode.usesDMSState() != mode.usesDMSState() || recorded != mode.usesDMSState() { + if (db.mode == ReadingDMS) != (mode == WritingDMS) || recorded != (mode == WritingDMS) { return errors.New("cannot change the deployment state source when upgrading to write") } - if bestEffort && !mode.isFileWithHistory() { + if bestEffort && mode != WritingFileWithHistory && mode != WritingFileWithHistoryEnforced { return errors.New("deployment_history_best_effort state requires writing with deployment history enabled") } @@ -1026,7 +1029,7 @@ func (db *DeploymentState) unlockedSave() error { // Data to persist in the remote resources.json file. func (db *DeploymentState) dataForFile() Database { - if db.isDeploymentMetadataService() { + if db.mode == ReadingDMS || db.mode == WritingDMS { header := db.Data.Header if header.Features == nil { header.Features = make(map[string]struct{}) From 399b1735815c836a3bdb025574c7b50537871854 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 12:34:46 +0000 Subject: [PATCH 10/52] Remove redundant version-completion commentary --- bundle/direct/dstate/state.go | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index e83eef71414..48b2d675daf 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -223,12 +223,7 @@ func (db *DeploymentState) RecordingError() error { func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (bool, error) { db.mu.Lock() buf := db.operationBuffer - // WritingFile has no DMS buffer, so version completion is skipped. - if buf == nil { - db.mu.Unlock() - return false, nil - } - if db.versionCompleted { + if buf == nil || db.versionCompleted { db.mu.Unlock() return false, nil } From 72fe431c5be4599c35a6bb896971641a2cffeab7 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 12:41:38 +0000 Subject: [PATCH 11/52] Make state save and delete modes explicit --- bundle/direct/dstate/state.go | 72 ++++++++++++++++++++++++++--------- 1 file changed, 54 insertions(+), 18 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 48b2d675daf..9166794a3d3 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -321,20 +321,40 @@ func (db *DeploymentState) SaveState(ctx context.Context, key, newID string, sta DependsOn: dependsOn, } - if db.mode == WritingDMS { + switch db.mode { + case WritingDMS: db.Data.State[key] = entry - } else if err := appendJSONLine(db.walFile, WALEntry{Key: key, Value: &entry}); err != nil { - return err - } - db.stateIDs[key] = newID + db.stateIDs[key] = newID + if buf := db.operationBuffer; buf != nil { + recorded, err := json.Marshal(RecordedState{State: entry.State, DependsOn: dependsOn}) + if err != nil { + return err + } + buf.RecordOperation(ctx, key, false, newID, recorded) + } - // File-backed modes record history only after the WAL append succeeds. - if buf := db.operationBuffer; buf != nil { - recorded, err := json.Marshal(RecordedState{State: entry.State, DependsOn: dependsOn}) - if err != nil { + case WritingFile: + if err := appendJSONLine(db.walFile, WALEntry{Key: key, Value: &entry}); err != nil { return err } - buf.RecordOperation(ctx, key, false, newID, recorded) + db.stateIDs[key] = newID + + case WritingFileWithHistory, WritingFileWithHistoryEnforced: + // File-backed modes record history only after the WAL append succeeds. + if err := appendJSONLine(db.walFile, WALEntry{Key: key, Value: &entry}); err != nil { + return err + } + db.stateIDs[key] = newID + if buf := db.operationBuffer; buf != nil { + recorded, err := json.Marshal(RecordedState{State: entry.State, DependsOn: dependsOn}) + if err != nil { + return err + } + buf.RecordOperation(ctx, key, false, newID, recorded) + } + + default: + return fmt.Errorf("cannot save state in mode %d", db.mode) } return nil } @@ -353,16 +373,32 @@ func (db *DeploymentState) DeleteState(ctx context.Context, key string, inProgre // Read before the delete below: DMS needs the id to say which resource went away. deletedID := db.stateIDs[key] - if db.mode == WritingDMS { + switch db.mode { + case WritingDMS: delete(db.Data.State, key) - } else if err := appendJSONLine(db.walFile, WALEntry{Key: key}); err != nil { - return err - } - if buf := db.operationBuffer; buf != nil { - buf.RecordOperation(ctx, key, inProgress, deletedID, nil) - } + if buf := db.operationBuffer; buf != nil { + buf.RecordOperation(ctx, key, inProgress, deletedID, nil) + } + delete(db.stateIDs, key) - delete(db.stateIDs, key) + case WritingFile: + if err := appendJSONLine(db.walFile, WALEntry{Key: key}); err != nil { + return err + } + delete(db.stateIDs, key) + + case WritingFileWithHistory, WritingFileWithHistoryEnforced: + if err := appendJSONLine(db.walFile, WALEntry{Key: key}); err != nil { + return err + } + if buf := db.operationBuffer; buf != nil { + buf.RecordOperation(ctx, key, inProgress, deletedID, nil) + } + delete(db.stateIDs, key) + + default: + return fmt.Errorf("cannot delete state in mode %d", db.mode) + } return nil } From a2157260883cfa21ddccaf2c5d4c796ea98b18bd Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 12:51:50 +0000 Subject: [PATCH 12/52] Guard version completion with RecordsHistory --- bundle/direct/dstate/state.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 9166794a3d3..e79bc93c34d 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -221,6 +221,10 @@ func (db *DeploymentState) RecordingError() error { // one does nothing. Finalize has already drained the buffered operations. It reads only fields that // survive Finalize's reset and asserts nothing, so it is safe to call after the state is closed. func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (bool, error) { + if !db.RecordsHistory() { + return false, nil + } + db.mu.Lock() buf := db.operationBuffer if buf == nil || db.versionCompleted { From 3042a1541109aab898f6dcc77ca9b10264f5a07d Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 12:54:52 +0000 Subject: [PATCH 13/52] Restore WithRecovery documentation --- bundle/direct/dstate/state.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index e79bc93c34d..904b5dfc4f0 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -503,7 +503,7 @@ func (db *DeploymentState) GetOrInitLineage() string { return db.Data.Lineage } -// WithRecovery allows Open to merge a leftover WAL into file state. +// If true, then Open reads the WAL and merges it in the state. If false, and WAL is present, Open returns an error. type WithRecovery bool // Open loads resource state using mode, recovering a leftover WAL when requested. From c6a030a29f90cf72ddbff4f6540be49ea8999fed Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 13:47:24 +0000 Subject: [PATCH 14/52] Use an unbounded DMS queue and enforce timeouts only at drain --- bundle/direct/dstate/mode_test.go | 66 ++++++++++----------------- bundle/direct/dstate/state.go | 3 -- libs/dms/operation_buffer.go | 76 ++++++++++++++++++++++--------- libs/dms/operation_buffer_test.go | 72 +++++++++++------------------ 4 files changed, 104 insertions(+), 113 deletions(-) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index 647d9f3d3b0..7df54de2302 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -48,8 +48,8 @@ func modeContext(t *testing.T, handler http.HandlerFunc) context.Context { client, err := databricks.NewWorkspaceClient(&databricks.Config{ Host: "https://state.test", Token: "test-token", AuthType: "pat", HTTPTransport: modeTransport{handler: handler}, - // Let our request deadline expire before the SDK's inactivity timeout. - HTTPTimeoutSeconds: 120, + // Keep SDK inactivity timeouts outside the drain budgets exercised here. + HTTPTimeoutSeconds: 600, }) require.NoError(t, err) return cmdctx.SetWorkspaceClient(t.Context(), client) @@ -319,42 +319,7 @@ func TestHistoryDrainTimeoutPreservesFileState(t *testing.T) { } } -func TestOperationTimeout(t *testing.T) { - for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { - t.Run(fmt.Sprint(mode), func(t *testing.T) { - synctest.Test(t, func(t *testing.T) { - finished := make(chan struct{}) - ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { - _, _ = io.Copy(io.Discard, r.Body) - <-r.Context().Done() - close(finished) - }) - path := filepath.Join(t.TempDir(), "resources.json") - var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) - db.InitializeOperationBuffer(ctx, "dep-1", 1, 100) - start := time.Now() - require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) - <-finished - synctest.Wait() - assert.Equal(t, time.Minute, time.Since(start)) - _, err := db.Finalize(ctx) - if mode == dstate.WritingDMS { - require.ErrorIs(t, err, context.DeadlineExceeded) - } else { - require.NoError(t, err) - var reopened dstate.DeploymentState - require.NoError(t, reopened.Open(ctx, path, dstate.ReadingFile, false, dstate.OpenDmsArgs{})) - assert.Equal(t, "id-1", reopened.GetResourceID("resources.jobs.job")) - _, err = reopened.Finalize(ctx) - require.NoError(t, err) - } - }) - }) - } -} - -func TestCompleteVersionTimeout(t *testing.T) { +func TestCompleteVersionHonorsCallerDeadline(t *testing.T) { for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { t.Run(fmt.Sprint(mode), func(t *testing.T) { synctest.Test(t, func(t *testing.T) { @@ -374,9 +339,11 @@ func TestCompleteVersionTimeout(t *testing.T) { require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) _, err := db.Finalize(ctx) require.NoError(t, err) + completionCtx, cancel := context.WithTimeout(ctx, 90*time.Second) + defer cancel() start := time.Now() - completed, err := db.CompleteVersion(ctx, true) - assert.Equal(t, time.Minute, time.Since(start)) + completed, err := db.CompleteVersion(completionCtx, true) + assert.Equal(t, 90*time.Second, time.Since(start)) assert.False(t, completed) if mode == dstate.WritingDMS { require.ErrorIs(t, err, context.DeadlineExceeded) @@ -508,13 +475,17 @@ func TestRecoveryRejectsUnknownWALFeatures(t *testing.T) { assert.FileExists(t, path+".wal") } -func TestFullHistoryBufferUnblocksAfterOperationTimeout(t *testing.T) { +func TestHistoryQueueOnlyTimesOutAtDrain(t *testing.T) { for _, mode := range []dstate.Mode{dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { t.Run(fmt.Sprint(mode), func(t *testing.T) { synctest.Test(t, func(t *testing.T) { + started := make(chan struct{}) + finished := make(chan struct{}) ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { _, _ = io.Copy(io.Discard, r.Body) + close(started) <-r.Context().Done() + close(finished) }) path := filepath.Join(t.TempDir(), "resources.json") var db dstate.DeploymentState @@ -526,12 +497,23 @@ func TestFullHistoryBufferUnblocksAfterOperationTimeout(t *testing.T) { key := fmt.Sprintf("resources.jobs.job_%d", i) require.NoError(t, db.SaveState(ctx, key, strconv.Itoa(i), struct{}{}, nil)) } - assert.Equal(t, time.Minute, time.Since(start)) + assert.Zero(t, time.Since(start)) + <-started + time.Sleep(time.Minute + time.Second) + select { + case <-finished: + t.Fatal("recording stopped before drain") + default: + } + drainStart := time.Now() exported, err := db.Finalize(ctx) require.NoError(t, err) + assert.Equal(t, 160*time.Second, time.Since(drainStart)) assert.Len(t, exported, operationCount) assert.FileExists(t, path) assert.NoFileExists(t, path+".wal") + <-finished + assert.Equal(t, 160*time.Second, time.Since(drainStart)) }) }) } diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 904b5dfc4f0..f9d521daaec 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -236,9 +236,6 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b bestEffort := db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced db.mu.Unlock() - ctx, cancel := context.WithTimeout(ctx, dms.OperationTimeout) - defer cancel() - // A recording failure fails the version even when the caller counted the deploy a success: the // service does not then hold everything the WAL does. Finalize already drained and surfaced it; // this reads the drained buffer's error so a destroy whose uploads failed keeps its record. diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index d3d7e7de81b..f1e419c58c0 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -5,19 +5,11 @@ import ( "encoding/json" "fmt" "sync" - "time" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/databricks-sdk-go/service/bundledeployments" ) -// bufferedOperations caps how far ahead of the service a deploy may get; DMS is what the next -// plan reads. -const bufferedOperations = 10 - -// OperationTimeout caps a single DMS history request, including its retries. -const OperationTimeout = time.Minute - // BufferOptions controls whether recording continues after an error. type BufferOptions struct { BestEffort bool @@ -30,11 +22,8 @@ type OperationBuffer struct { deploymentID string versionNum int - // The buffer. queue holds bundle state keys, and pending the newest update per key, so a - // second write for a resource replaces the first. - queue chan string + wake chan struct{} done chan struct{} - stopQueue func() cancel context.CancelFunc bestEffort bool @@ -50,6 +39,10 @@ type OperationBuffer struct { // mu guards the fields below. mu sync.Mutex + // queue is unbounded and contains one key per pending update. + queue []string + closed bool + // pending holds the newest update per resource key, absent once the writer takes it. pending map[string]OperationUpdate @@ -69,7 +62,7 @@ func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum i b := &OperationBuffer{ deploymentID: deploymentID, versionNum: versionNum, - queue: make(chan string, bufferedOperations), + wake: make(chan struct{}, 1), done: make(chan struct{}), cancel: cancel, bestEffort: options.BestEffort, @@ -77,7 +70,6 @@ func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum i latestState: make(map[string]OperationUpdate), sequenceIDs: make(map[string]int64), } - b.stopQueue = sync.OnceFunc(func() { close(b.queue) }) go b.run(ctx) return b } @@ -107,10 +99,13 @@ func (b *OperationBuffer) RecordFailure(resourceKey, resourceID string, state js // that latest state. latestState and pending are updated under the one lock. Panics after Drain. func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, keepExistingUpdateIfAny bool, cause error) { b.mu.Lock() + defer b.mu.Unlock() if b.bestEffort && b.err != nil { - b.mu.Unlock() return } + if b.closed { + panic("recording operation after drain") + } if keepExistingUpdateIfAny { if latest, recorded := b.latestState[resourceKey]; recorded { update = latest.Merge(failureUpdate(cause)) @@ -123,16 +118,48 @@ func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, kee update = waiting.Merge(update) } b.pending[resourceKey] = update - b.mu.Unlock() // when already queued: the writer reads the map when it gets to the key, // so it picks up what was just stored. // // Only when the resource is not already queued, enqueue it. if !queued { + b.queue = append(b.queue, resourceKey) + b.notify() + } +} + +// notify wakes the writer without blocking producers. +func (b *OperationBuffer) notify() { + select { + case b.wake <- struct{}{}: + default: + } +} + +// next waits for a queued key, drain, or cancellation. +func (b *OperationBuffer) next(ctx context.Context) (string, bool) { + for { + b.mu.Lock() + if len(b.queue) > 0 { + key := b.queue[0] + b.queue[0] = "" + b.queue = b.queue[1:] + if len(b.queue) == 0 { + b.queue = nil + } + b.mu.Unlock() + return key, true + } + closed := b.closed + b.mu.Unlock() + if closed { + return "", false + } select { - case b.queue <- resourceKey: - case <-b.done: + case <-b.wake: + case <-ctx.Done(): + return "", false } } } @@ -151,11 +178,15 @@ func (b *OperationBuffer) run(ctx context.Context) { defer close(b.done) defer b.cancel() - for resourceKey := range b.queue { + for { + resourceKey, ok := b.next(ctx) if err := ctx.Err(); err != nil { b.setErr(err) return } + if !ok { + return + } update, ok := b.take(resourceKey) if !ok { // Unreachable: a key is queued only when nothing was waiting for it. Guard so @@ -175,8 +206,6 @@ func (b *OperationBuffer) run(ctx context.Context) { // write sends one update, at the sequence id the resource is at. func (b *OperationBuffer) write(ctx context.Context, key string, update OperationUpdate) error { - ctx, cancel := context.WithTimeout(ctx, OperationTimeout) - defer cancel() w := cmdctx.WorkspaceClient(ctx) result, err := w.BundleDeployments.UpdateOperation(ctx, bundledeployments.UpdateOperationRequest{ Name: OperationName(b.deploymentID, b.versionNum, key), @@ -195,7 +224,10 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio // Drain waits for buffered writes, canceling outstanding requests when ctx expires. // It returns the first recording error and is safe to call twice. func (b *OperationBuffer) Drain(ctx context.Context) error { - b.stopQueue() + b.mu.Lock() + b.closed = true + b.mu.Unlock() + b.notify() select { case <-b.done: return b.Err() diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index bb7259602df..117e2d4707e 100644 --- a/libs/dms/operation_buffer_test.go +++ b/libs/dms/operation_buffer_test.go @@ -3,8 +3,8 @@ package dms import ( "encoding/json" "errors" + "fmt" "strings" - "sync" "testing" "testing/synctest" @@ -21,7 +21,7 @@ import ( // queued builds a buffer whose queue nothing drains, so a test drives record and take itself. func queued() *OperationBuffer { return &OperationBuffer{ - queue: make(chan string, bufferedOperations), + wake: make(chan struct{}, 1), cancel: func() {}, pending: make(map[string]OperationUpdate), latestState: make(map[string]OperationUpdate), @@ -71,15 +71,10 @@ func TestBufferFailsOnOversizedState(t *testing.T) { } func TestDrainIsIdempotent(t *testing.T) { - // Nothing recorded, so no transport is needed: this is the second drain, which must not - // panic on an already closed queue. - s := queued() - s.done = make(chan struct{}) - s.stopQueue = sync.OnceFunc(func() { close(s.queue) }) - go s.run(t.Context()) - - require.NoError(t, s.Drain(t.Context())) - require.NoError(t, s.Drain(t.Context())) + // An idle writer must wake up when the queue is closed. + b := StartOperationBuffer(t.Context(), "dep-1", 1, BufferOptions{}) + require.NoError(t, b.Drain(t.Context())) + require.NoError(t, b.Drain(t.Context())) } func TestRecordFailureKeepsTheStateTheLastWriteRecorded(t *testing.T) { @@ -117,52 +112,37 @@ func TestRecordFailureUsesTheGivenStateWhenNoWriteRecorded(t *testing.T) { assert.Equal(t, bundledeployments.OperationStatusOperationStatusFailed, update.Status) } -func TestBestEffortBufferWaitsForQueueSpace(t *testing.T) { +func TestBufferQueuesWithoutBackpressure(t *testing.T) { synctest.Test(t, func(t *testing.T) { b := queued() - b.bestEffort = true - for range bufferedOperations { - b.queue <- "occupied" - } - done := make(chan struct{}) - go func() { - b.RecordOperation(t.Context(), "resources.jobs.foo", false, "id-1", json.RawMessage(`{}`)) - close(done) - }() - synctest.Wait() - select { - case <-done: - t.Fatal("recording should wait for queue space") - default: + const operationCount = 1000 + for i := range operationCount { + key := fmt.Sprintf("resources.jobs.job_%d", i) + b.RecordOperation(t.Context(), key, false, "id-1", json.RawMessage(`{"state":{}}`)) } - <-b.queue - <-done assert.NoError(t, b.Err()) - assert.Len(t, b.queue, bufferedOperations) - _, ok := b.take("resources.jobs.foo") - assert.True(t, ok) + assert.Len(t, b.queue, operationCount) + assert.Len(t, b.pending, operationCount) + for i := range operationCount { + key, ok := b.next(t.Context()) + require.True(t, ok) + assert.Equal(t, fmt.Sprintf("resources.jobs.job_%d", i), key) + _, ok = b.take(key) + require.True(t, ok) + } + assert.Empty(t, b.queue) + assert.Empty(t, b.pending) }) } -func TestBestEffortBufferUnblocksWhenWriterStops(t *testing.T) { +func TestBestEffortBufferStopsIdleWriterAfterError(t *testing.T) { synctest.Test(t, func(t *testing.T) { - b := queued() - b.bestEffort = true - b.done = make(chan struct{}) - for range bufferedOperations { - b.queue <- "occupied" - } - done := make(chan struct{}) - go func() { - b.RecordOperation(t.Context(), "resources.jobs.foo", false, "id-1", json.RawMessage(`{}`)) - close(done) - }() + b := StartOperationBuffer(t.Context(), "dep-1", 1, BufferOptions{BestEffort: true}) synctest.Wait() cause := errors.New("recording failed") b.setErr(cause) - close(b.done) - <-done - assert.ErrorIs(t, b.Err(), cause) + <-b.done + assert.ErrorIs(t, b.Drain(t.Context()), cause) }) } From 85cb14170ff8d3936eadf7af56f97110e49a8e43 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 15:43:36 +0000 Subject: [PATCH 15/52] Remove redundant mode validation and clarify history errors --- bundle/direct/dstate/mode.go | 11 ----------- bundle/direct/dstate/mode_test.go | 3 +-- bundle/direct/dstate/state.go | 7 +++---- 3 files changed, 4 insertions(+), 17 deletions(-) diff --git a/bundle/direct/dstate/mode.go b/bundle/direct/dstate/mode.go index 5280af1db39..72592f65f49 100644 --- a/bundle/direct/dstate/mode.go +++ b/bundle/direct/dstate/mode.go @@ -1,7 +1,5 @@ package dstate -import "fmt" - // Mode selects the source of resource state and the persistence required for writes. type Mode int @@ -40,12 +38,3 @@ func (m Mode) isWritable() bool { return false } } - -func (m Mode) validate() error { - switch m { - case ReadingFile, WritingFile, ReadingDMS, WritingDMS, WritingFileWithHistory, WritingFileWithHistoryEnforced: - return nil - default: - return fmt.Errorf("invalid deployment state mode: %d", m) - } -} diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index 7df54de2302..f85a7e01c97 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -388,14 +388,13 @@ func TestEnforcedHistorySurvivesWALRecovery(t *testing.T) { } } -func TestOpenRejectsInvalidModesAndFeatures(t *testing.T) { +func TestOpenRejectsIncompatibleModesAndFeatures(t *testing.T) { for _, tc := range []struct { name string mode dstate.Mode features string wantErr string }{ - {"unknown mode", dstate.Mode(99), `{}`, "invalid deployment state mode"}, {"unknown feature", dstate.ReadingFile, `{"future":{}}`, "requires features"}, {"conflicting features", dstate.ReadingFile, `{"deployment_history":{},"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, {"DMS reading best effort history", dstate.ReadingDMS, `{"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index f9d521daaec..3b1bdb0e228 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -203,9 +203,11 @@ func (db *DeploymentState) getOperationBuffer() *dms.OperationBuffer { // RecordingError reports whether recording state writes to the service has failed, so the apply // stops touching resources once the service is no longer keeping up. File-backed history modes -// always return nil because their resource state remains recoverable from the WAL. +// always return nil so history recording errors do not block the deployment; resource state +// remains recoverable from the WAL. func (db *DeploymentState) RecordingError() error { if !db.IsDeploymentMetadataService() { + // TODO: Warn users when history recording fails in file-backed history modes. return nil } buf := db.getOperationBuffer() @@ -536,9 +538,6 @@ func (db *DeploymentState) reset() { } func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode Mode, withRecovery WithRecovery, dmsDeployment OpenDmsArgs) error { - if err := mode.validate(); err != nil { - return err - } db.Path = path db.mode = ReadingFile if mode == ReadingDMS || mode == WritingDMS { From 7818b7ae982947ae36d240785ec9e25819af5e0c Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 15:48:58 +0000 Subject: [PATCH 16/52] Document Open behavior for each deployment state mode --- bundle/direct/dstate/state.go | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 3b1bdb0e228..cc87be9b633 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -505,8 +505,24 @@ func (db *DeploymentState) GetOrInitLineage() string { // If true, then Open reads the WAL and merges it in the state. If false, and WAL is present, Open returns an error. type WithRecovery bool -// Open loads resource state using mode, recovering a leftover WAL when requested. -// DMS reads and history writes require a workspace client on ctx. +// Open loads deployment state and enables reads or writes according to mode: +// - ReadingFile: reads resources.json; state writes are disabled. +// - WritingFile: reads resources.json and opens a WAL for state writes. +// - ReadingDMS: reads resources from DMS; state writes are disabled. +// - WritingDMS: reads and writes resources through DMS; no WAL is opened. +// - WritingFileWithHistory: behaves like WritingFile and also records DMS history. +// History errors do not fail the deployment. +// - WritingFileWithHistoryEnforced: also persists a feature marker that rejects +// older clients that do not support history recording. +// +// All modes read deployment metadata and feature markers from resources.json. +// Open rejects modes that conflict with those markers. +// +// File modes recover a leftover WAL only when withRecovery is set. +// DMS modes reject any leftover WAL. +// +// DMS reads require a workspace client on ctx. DMS writes and history recording +// start after a version exists and InitializeOperationBuffer is called. func (db *DeploymentState) Open(ctx context.Context, path string, mode Mode, withRecovery WithRecovery, dmsDeployment OpenDmsArgs) error { db.mu.Lock() defer db.mu.Unlock() From 0f7af033cfce95fac69e30f24ff45215b7607ae2 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 15:52:55 +0000 Subject: [PATCH 17/52] Separate write initialization from mode transitions --- bundle/direct/dstate/state.go | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index cc87be9b633..41b96ed38ad 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -555,10 +555,7 @@ func (db *DeploymentState) reset() { func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode Mode, withRecovery WithRecovery, dmsDeployment OpenDmsArgs) error { db.Path = path - db.mode = ReadingFile - if mode == ReadingDMS || mode == WritingDMS { - db.mode = ReadingDMS - } + db.mode = mode db.operationBuffer = nil db.versionCompleted = false db.historyDrainTimeout = 0 @@ -676,7 +673,7 @@ Run "databricks bundle destroy" first, then deploy again with deployment history } if mode.isWritable() { - return db.unlockedUpgradeToWrite(mode) + return db.initializeWrites(mode) } return nil @@ -898,10 +895,6 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported func (db *DeploymentState) UpgradeToWrite(mode Mode) error { db.mu.Lock() defer db.mu.Unlock() - return db.unlockedUpgradeToWrite(mode) -} - -func (db *DeploymentState) unlockedUpgradeToWrite(mode Mode) error { if db.Path == "" { return errors.New("internal error: DeploymentState must be opened first") } @@ -911,12 +904,25 @@ func (db *DeploymentState) unlockedUpgradeToWrite(mode Mode) error { if !mode.isWritable() { return fmt.Errorf("invalid deployment state write mode: %d", mode) } + if (db.mode == ReadingDMS) != (mode == WritingDMS) { + return errors.New("cannot change the deployment state source when upgrading to write") + } + if err := db.initializeWrites(mode); err != nil { + return err + } + db.mode = mode + return nil +} + +// initializeWrites validates write compatibility and prepares DMS metadata or a WAL. +// Callers must hold db.mu. +func (db *DeploymentState) initializeWrites(mode Mode) error { if err := assertNoUnsupportedFeatures(db.Data.Features); err != nil { return err } _, recorded := db.Data.Features[FeatureDeploymentHistory] _, bestEffort := db.Data.Features[FeatureDeploymentHistoryBestEffort] - if (db.mode == ReadingDMS) != (mode == WritingDMS) || recorded != (mode == WritingDMS) { + if recorded != (mode == WritingDMS) { return errors.New("cannot change the deployment state source when upgrading to write") } if bestEffort && mode != WritingFileWithHistory && mode != WritingFileWithHistoryEnforced { @@ -927,7 +933,6 @@ func (db *DeploymentState) unlockedUpgradeToWrite(mode Mode) error { db.Data.Lineage = db.GetOrInitLineage() db.Data.StateVersion = currentStateVersion db.Data.CLIVersion = build.GetInfo().Version - db.mode = mode return nil } @@ -960,7 +965,6 @@ func (db *DeploymentState) unlockedUpgradeToWrite(mode Mode) error { return fmt.Errorf("failed to write WAL header to %s: %w", walPath, err) } db.walFile = walFile - db.mode = mode return nil } From 58b054989de0960c9333aff198868c7ad1e72ae1 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 15:55:06 +0000 Subject: [PATCH 18/52] Restore DMS helper call for serial lookup --- bundle/direct/dstate/state.go | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 41b96ed38ad..ea91f7a3052 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -455,7 +455,7 @@ func (db *DeploymentState) StateFeatures() map[string]struct{} { func (db *DeploymentState) GetSerial() int { db.mu.Lock() defer db.mu.Unlock() - if db.mode == ReadingDMS || db.mode == WritingDMS { + if db.isDeploymentMetadataService() { return db.VersionID } return db.Data.Serial @@ -467,6 +467,10 @@ func (db *DeploymentState) GetSerial() int { func (db *DeploymentState) IsDeploymentMetadataService() bool { db.mu.Lock() defer db.mu.Unlock() + return db.isDeploymentMetadataService() +} + +func (db *DeploymentState) isDeploymentMetadataService() bool { return db.mode == ReadingDMS || db.mode == WritingDMS } From 09d854cfe3f94fad57f622cdcc2a7515b70fe2aa Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 15:58:31 +0000 Subject: [PATCH 19/52] Use both version counters for file-backed deployment history --- bundle/direct/dstate/mode_test.go | 28 ++++++++++++++++++++++++++++ bundle/direct/dstate/state.go | 14 +++++++++----- 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index f85a7e01c97..9b3f082ef23 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -55,6 +55,34 @@ func modeContext(t *testing.T, handler http.HandlerFunc) context.Context { return cmdctx.SetWorkspaceClient(t.Context(), client) } +func TestGetSerialWithDivergingCounters(t *testing.T) { + for _, tc := range []struct { + mode dstate.Mode + versionAhead int + serialAhead int + }{ + {dstate.ReadingFile, 4, 7}, + {dstate.WritingFile, 4, 7}, + {dstate.ReadingDMS, 7, 4}, + {dstate.WritingDMS, 7, 4}, + {dstate.WritingFileWithHistory, 7, 7}, + {dstate.WritingFileWithHistoryEnforced, 7, 7}, + } { + t.Run(fmt.Sprint(tc.mode), func(t *testing.T) { + ctx := modeContext(t, nil) + var db dstate.DeploymentState + path := filepath.Join(t.TempDir(), "resources.json") + require.NoError(t, db.Open(ctx, path, tc.mode, false, dstate.OpenDmsArgs{})) + db.Data.Serial, db.VersionID = 4, 7 + assert.Equal(t, tc.versionAhead, db.GetSerial(), "DMS version ahead of file serial") + db.Data.Serial, db.VersionID = 7, 4 + assert.Equal(t, tc.serialAhead, db.GetSerial(), "file serial ahead of DMS version") + _, err := db.Finalize(ctx) + require.NoError(t, err) + }) + } +} + func TestModeAssertions(t *testing.T) { for _, tc := range []struct { mode dstate.Mode diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index ea91f7a3052..76a0291d01a 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -449,16 +449,20 @@ func (db *DeploymentState) StateFeatures() map[string]struct{} { return maps.Clone(db.Data.Features) } -// GetSerial reports the counter a saved plan is validated against: the service's version for a -// deployment that records history, and the state serial otherwise. When the deployment uses the -// deployment metadata service, the serial from resources.json is not used for plan validation. +// GetSerial reports the counter used to validate saved plans. +// DMS modes use the service version; plain file modes use the file serial. +// File modes with history use the larger counter because either can advance independently. func (db *DeploymentState) GetSerial() int { db.mu.Lock() defer db.mu.Unlock() - if db.isDeploymentMetadataService() { + switch db.mode { + case ReadingDMS, WritingDMS: return db.VersionID + case WritingFileWithHistory, WritingFileWithHistoryEnforced: + return max(db.VersionID, db.Data.Serial) + default: + return db.Data.Serial } - return db.Data.Serial } // IsDeploymentMetadataService reports whether this deployment's resource state lives in the From 09161b19f8129502512737a8c01bc9ace640dd0b Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Thu, 22 Oct 2026 16:13:53 +0000 Subject: [PATCH 20/52] Clear deployment state during reset and version completion --- bundle/direct/dstate/mode_test.go | 36 +++++++++++++++++++------------ bundle/direct/dstate/state.go | 21 +++++++++--------- bundle/phases/deploy.go | 1 + bundle/phases/destroy.go | 4 ++-- bundle/phases/telemetry.go | 7 +++++- 5 files changed, 42 insertions(+), 27 deletions(-) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index 9b3f082ef23..38a9884670c 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -373,6 +373,9 @@ func TestCompleteVersionHonorsCallerDeadline(t *testing.T) { completed, err := db.CompleteVersion(completionCtx, true) assert.Equal(t, 90*time.Second, time.Since(start)) assert.False(t, completed) + assert.False(t, db.RecordsHistory()) + assert.Empty(t, db.DeploymentID) + assert.Zero(t, db.VersionID) if mode == dstate.WritingDMS { require.ErrorIs(t, err, context.DeadlineExceeded) } else { @@ -459,21 +462,26 @@ func TestCompleteVersionAfterFinalize(t *testing.T) { _, _ = w.Write([]byte(`{"sequence_id":1}`)) } }) - path := filepath.Join(t.TempDir(), "resources.json") var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) - db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) - require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) - _, err := db.Finalize(ctx) - require.NoError(t, err) - assert.False(t, db.IsOpen()) - completed, err := db.CompleteVersion(ctx, true) - require.NoError(t, err) - assert.True(t, completed) - completed, err = db.CompleteVersion(ctx, true) - require.NoError(t, err) - assert.False(t, completed) - assert.Equal(t, 1, completions) + for cycle := range 2 { + path := filepath.Join(t.TempDir(), "resources.json") + require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) + db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) + require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) + _, err := db.Finalize(ctx) + require.NoError(t, err) + assert.False(t, db.IsOpen()) + completed, err := db.CompleteVersion(ctx, true) + require.NoError(t, err) + assert.True(t, completed) + completed, err = db.CompleteVersion(ctx, true) + require.NoError(t, err) + assert.False(t, completed) + assert.Equal(t, cycle+1, completions) + assert.False(t, db.RecordsHistory()) + assert.Empty(t, db.DeploymentID) + assert.Zero(t, db.VersionID) + } }) } } diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 76a0291d01a..0f85aa4a302 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -236,6 +236,9 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b db.versionCompleted = true deploymentID, versionID := db.DeploymentID, db.VersionID bestEffort := db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced + if db.Path == "" { + db.reset() + } db.mu.Unlock() // A recording failure fails the version even when the caller counted the deploy a success: the @@ -543,8 +546,6 @@ func (db *DeploymentState) Open(ctx context.Context, path string, mode Mode, wit if err != nil { // A failed open must leave the receiver closed so it can be retried. db.reset() - db.mode = ReadingFile - db.operationBuffer = nil } return err } @@ -559,14 +560,19 @@ func (db *DeploymentState) reset() { db.Data = Database{} db.stateIDs = nil db.historyDrainTimeout = 0 + // Retain recording state only while version completion is pending. + if db.operationBuffer == nil || db.versionCompleted { + db.mode = ReadingFile + db.operationBuffer = nil + db.versionCompleted = false + db.DeploymentID = "" + db.VersionID = 0 + } } func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode Mode, withRecovery WithRecovery, dmsDeployment OpenDmsArgs) error { db.Path = path db.mode = mode - db.operationBuffer = nil - db.versionCompleted = false - db.historyDrainTimeout = 0 db.DeploymentID = dmsDeployment.DeploymentID db.VersionID = dmsDeployment.LastVersionID @@ -699,11 +705,6 @@ func (db *DeploymentState) OpenWithData(path string, data Database) { db.Path = path db.mode = ReadingFile - db.operationBuffer = nil - db.versionCompleted = false - db.historyDrainTimeout = 0 - db.DeploymentID = "" - db.VersionID = 0 db.Data = data db.stateIDs = make(map[string]string) for key, entry := range data.State { diff --git a/bundle/phases/deploy.go b/bundle/phases/deploy.go index 18e4ba4f4ef..b906fa70bfe 100644 --- a/bundle/phases/deploy.go +++ b/bundle/phases/deploy.go @@ -156,6 +156,7 @@ func uploadLibraries(ctx context.Context, b *bundle.Bundle, libs map[string][]li // If readPlanPath is provided, the plan is loaded from that file instead of being calculated. func Deploy(ctx context.Context, b *bundle.Bundle, outputHandler sync.OutputHandler, libs map[string][]libraries.LocationToUpdate, plan *deployplan.Plan, dmsDeployment *bundledeployments.Deployment) { log.Info(ctx, "Phase: deploy") + b.Metrics.SetBoolValue(metrics.DeploymentHistoryEnabled, b.DeploymentBundle.StateDB.IsDeploymentMetadataService()) // Core mutators that CRUD resources and modify deployment state. These // mutators need informed consent if they are potentially destructive. The diff --git a/bundle/phases/destroy.go b/bundle/phases/destroy.go index 7b042a456b2..6d4a1e6e049 100644 --- a/bundle/phases/destroy.go +++ b/bundle/phases/destroy.go @@ -159,6 +159,7 @@ func destroyCore(ctx context.Context, b *bundle.Bundle, plan *deployplan.Plan) { if b.DeploymentBundle.StateDB.IsDeploymentMetadataService() { // Complete version before deleting remote files; the deployment node is under statePath. + deploymentID := b.DeploymentBundle.StateDB.DeploymentID completed, err := b.DeploymentBundle.StateDB.CompleteVersion(ctx, true) if err != nil { logdiag.LogError(ctx, err) @@ -166,7 +167,6 @@ func destroyCore(ctx context.Context, b *bundle.Bundle, plan *deployplan.Plan) { } // A completed destroy's resources are gone, so its deployment record is deleted too. if completed { - deploymentID := b.DeploymentBundle.StateDB.DeploymentID w := b.WorkspaceClient(ctx) if err := deleteDeployment(ctx, w, deploymentID); err != nil { logdiag.LogError(ctx, fmt.Errorf("failed to delete deployment: %w", err)) @@ -295,11 +295,11 @@ func Destroy(ctx context.Context, b *bundle.Bundle) { // destroyCore has completed the version. defer func() { if b.DeploymentBundle.StateDB.IsDeploymentMetadataService() { + deploymentID := b.DeploymentBundle.StateDB.DeploymentID completed, err := b.DeploymentBundle.StateDB.CompleteVersion(ctx, !logdiag.HasError(ctx)) if err != nil { logdiag.LogError(ctx, err) } else if completed { - deploymentID := b.DeploymentBundle.StateDB.DeploymentID w := b.WorkspaceClient(ctx) if err := deleteDeployment(ctx, w, deploymentID); err != nil { logdiag.LogError(ctx, fmt.Errorf("failed to delete deployment: %w", err)) diff --git a/bundle/phases/telemetry.go b/bundle/phases/telemetry.go index 99464f3ef24..013512b8aba 100644 --- a/bundle/phases/telemetry.go +++ b/bundle/phases/telemetry.go @@ -198,7 +198,12 @@ func aiRuntimeTaskMetrics(jobs map[string]*resources.Job) (present, scheduled, m // LogDeployTelemetry logs a telemetry event for a bundle deploy command. func LogDeployTelemetry(ctx context.Context, b *bundle.Bundle, errMsg string) { errMsg = telemetry.ScrubErrorMessage(errMsg) - b.Metrics.SetBoolValue(metrics.DeploymentHistoryEnabled, b.DeploymentBundle.StateDB.IsDeploymentMetadataService()) + // Deploy captures this before state cleanup. Fall back to the current state for earlier failures. + if !slices.ContainsFunc(b.Metrics.BoolValues, func(v protos.BoolMapEntry) bool { + return v.Key == metrics.DeploymentHistoryEnabled + }) { + b.Metrics.SetBoolValue(metrics.DeploymentHistoryEnabled, b.DeploymentBundle.StateDB.IsDeploymentMetadataService()) + } resourcesCount := int64(0) _, err := dyn.MapByPattern(b.Config.Value(), dyn.NewPattern(dyn.Key("resources"), dyn.AnyKey(), dyn.AnyKey()), func(p dyn.Path, v dyn.Value) (dyn.Value, error) { From 5beb2cd76d9abbef255ffd017dd503756a6c5838 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:01:58 +0000 Subject: [PATCH 21/52] Simplify WAL feature recovery to use header features directly --- bundle/direct/dstate/mode_test.go | 11 ----------- bundle/direct/dstate/state.go | 11 +---------- 2 files changed, 1 insertion(+), 21 deletions(-) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index 38a9884670c..c77020704cd 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -499,17 +499,6 @@ func TestFailedUpgradeLeavesReadModeAvailable(t *testing.T) { require.NoError(t, err) } -func TestRecoveryRejectsUnknownWALFeatures(t *testing.T) { - path := filepath.Join(t.TempDir(), "resources.json") - wal := `{"state_version":3,"lineage":"lineage","serial":1,"features":{"future":{}}}` + "\n" + `{"k":"resources.jobs.job","v":{"__id__":"id-1","state":{}}}` + "\n" - require.NoError(t, os.WriteFile(path+".wal", []byte(wal), 0o600)) - var db dstate.DeploymentState - require.ErrorContains(t, db.Open(t.Context(), path, dstate.ReadingFile, true, dstate.OpenDmsArgs{}), "future") - assert.False(t, db.IsOpen()) - assert.NoFileExists(t, path) - assert.FileExists(t, path+".wal") -} - func TestHistoryQueueOnlyTimesOutAtDrain(t *testing.T) { for _, mode := range []dstate.Mode{dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { t.Run(fmt.Sprint(mode), func(t *testing.T) { diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 0f85aa4a302..e323adf9d9a 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -778,16 +778,7 @@ func (db *DeploymentState) mergeWalIntoState(ctx context.Context) (bool, error) if header.Serial > expectedSerial { return false, fmt.Errorf("WAL serial (%d) is ahead of expected (%d), state may be corrupted", header.Serial, expectedSerial) } - if err := assertNoUnsupportedFeatures(header.Features); err != nil { - return false, err - } - newFeatures = maps.Clone(db.Data.Features) - for feature := range header.Features { - if newFeatures == nil { - newFeatures = make(map[string]struct{}) - } - newFeatures[feature] = struct{}{} - } + newFeatures = header.Features newSerial = header.Serial newCLIVersion = header.CLIVersion } else { From 3dc54f15e34dc44a2b1eb3d56913bdb729a19387 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:04:21 +0000 Subject: [PATCH 22/52] Explain why WAL feature changes need persistence --- bundle/direct/dstate/state.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index e323adf9d9a..626b0f97614 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -815,6 +815,9 @@ func (db *DeploymentState) mergeWalIntoState(ctx context.Context) (bool, error) } } + // Persist resource entries after the header, or feature changes in the header itself. + // Enforced history can add its feature marker without any resource writes; that marker + // still needs to reach resources.json before the WAL is removed. hasEntries := lineNumber > 1 || lineNumber == 1 && !maps.Equal(db.Data.Features, newFeatures) // Only advance the serial when the WAL carried changes, because the caller From 535cda74eadfb44891bc0a3538014177756036e6 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:08:01 +0000 Subject: [PATCH 23/52] Explain state source validation when upgrading to write --- bundle/direct/dstate/state.go | 1 + 1 file changed, 1 insertion(+) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 626b0f97614..2ff252f9888 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -907,6 +907,7 @@ func (db *DeploymentState) UpgradeToWrite(mode Mode) error { if !mode.isWritable() { return fmt.Errorf("invalid deployment state write mode: %d", mode) } + // Keep the state source unchanged: DMS reads require DMS writes; file reads require file writes. if (db.mode == ReadingDMS) != (mode == WritingDMS) { return errors.New("cannot change the deployment state source when upgrading to write") } From c8c6be13368e4b53a6f803fae9395ca9c5be43b7 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:15:47 +0000 Subject: [PATCH 24/52] Use one writable DMS mode and skip DMS write upgrades --- bundle/direct/dstate/mode.go | 6 +--- bundle/direct/dstate/mode_test.go | 19 ++++++------ bundle/direct/dstate/state.go | 47 +++++++++++++----------------- bundle/direct/dstate/state_test.go | 4 +-- bundle/phases/deploy.go | 6 +--- bundle/phases/destroy.go | 6 +--- cmd/bundle/utils/process.go | 4 +-- 7 files changed, 38 insertions(+), 54 deletions(-) diff --git a/bundle/direct/dstate/mode.go b/bundle/direct/dstate/mode.go index 72592f65f49..c7d6440573c 100644 --- a/bundle/direct/dstate/mode.go +++ b/bundle/direct/dstate/mode.go @@ -12,11 +12,7 @@ const ( // Finalize merges them into resources.json. DMS is unused. WritingFile - // ReadingDMS reads resource state from DMS. - // State writes are disabled. - ReadingDMS - - // WritingDMS records resource state in DMS without creating a WAL. + // WritingDMS reads and writes resource state through DMS without creating a WAL. // DMS errors fail the deployment. WritingDMS diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index c77020704cd..b473891234d 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -63,7 +63,6 @@ func TestGetSerialWithDivergingCounters(t *testing.T) { }{ {dstate.ReadingFile, 4, 7}, {dstate.WritingFile, 4, 7}, - {dstate.ReadingDMS, 7, 4}, {dstate.WritingDMS, 7, 4}, {dstate.WritingFileWithHistory, 7, 7}, {dstate.WritingFileWithHistoryEnforced, 7, 7}, @@ -92,7 +91,6 @@ func TestModeAssertions(t *testing.T) { }{ {dstate.ReadingFile, false, false, false}, {dstate.WritingFile, true, false, false}, - {dstate.ReadingDMS, false, true, true}, {dstate.WritingDMS, true, true, true}, {dstate.WritingFileWithHistory, true, false, true}, {dstate.WritingFileWithHistoryEnforced, true, false, true}, @@ -109,7 +107,11 @@ func TestModeAssertions(t *testing.T) { assert.NotPanics(t, db.AssertOpenedForReadOrWrite) if tc.writable { assert.NotPanics(t, db.AssertOpenedForWrite) - assert.Panics(t, db.AssertOpenedForRead) + if tc.dms { + assert.NotPanics(t, db.AssertOpenedForRead) + } else { + assert.Panics(t, db.AssertOpenedForRead) + } } else { assert.NotPanics(t, db.AssertOpenedForRead) assert.Panics(t, db.AssertOpenedForWrite) @@ -195,7 +197,7 @@ func TestDMSReadsAndWritesWithoutWAL(t *testing.T) { }) path := filepath.Join(t.TempDir(), "resources.json") var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, dstate.ReadingDMS, false, dstate.OpenDmsArgs{DeploymentID: "dep-1", LastVersionID: 7})) + require.NoError(t, db.Open(ctx, path, dstate.WritingDMS, false, dstate.OpenDmsArgs{DeploymentID: "dep-1", LastVersionID: 7})) assert.Equal(t, "remote-id", db.GetResourceID("resources.jobs.remote")) assert.Equal(t, 7, db.GetSerial()) require.NoError(t, db.UpgradeToWrite(dstate.WritingDMS)) @@ -228,10 +230,9 @@ func TestUpgradeToWriteValidatesMode(t *testing.T) { {"file", dstate.ReadingFile, dstate.WritingFile, false, ""}, {"history", dstate.ReadingFile, dstate.WritingFileWithHistory, false, ""}, {"enforced history", dstate.ReadingFile, dstate.WritingFileWithHistoryEnforced, false, ""}, - {"DMS", dstate.ReadingDMS, dstate.WritingDMS, false, ""}, - {"file to DMS", dstate.ReadingFile, dstate.WritingDMS, false, "cannot change"}, - {"DMS to file", dstate.ReadingDMS, dstate.WritingFile, false, "cannot change"}, - {"DMS to history", dstate.ReadingDMS, dstate.WritingFileWithHistory, false, "cannot change"}, + {"DMS", dstate.WritingDMS, dstate.WritingDMS, false, ""}, + {"file to DMS", dstate.ReadingFile, dstate.WritingDMS, false, "invalid"}, + {"DMS needs no WAL upgrade", dstate.WritingDMS, dstate.WritingFile, false, ""}, {"reading mode", dstate.ReadingFile, dstate.ReadingFile, false, "invalid"}, {"unknown mode", dstate.ReadingFile, dstate.Mode(99), false, "invalid"}, {"enforcement retained", dstate.ReadingFile, dstate.WritingFileWithHistory, true, ""}, @@ -428,7 +429,7 @@ func TestOpenRejectsIncompatibleModesAndFeatures(t *testing.T) { }{ {"unknown feature", dstate.ReadingFile, `{"future":{}}`, "requires features"}, {"conflicting features", dstate.ReadingFile, `{"deployment_history":{},"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, - {"DMS reading best effort history", dstate.ReadingDMS, `{"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, + {"DMS with best effort history", dstate.WritingDMS, `{"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, {"file writing DMS", dstate.WritingFile, `{"deployment_history":{}}`, "not supported"}, {"file writing best effort history", dstate.WritingFile, `{"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, } { diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 2ff252f9888..4cb8a626646 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -459,7 +459,7 @@ func (db *DeploymentState) GetSerial() int { db.mu.Lock() defer db.mu.Unlock() switch db.mode { - case ReadingDMS, WritingDMS: + case WritingDMS: return db.VersionID case WritingFileWithHistory, WritingFileWithHistoryEnforced: return max(db.VersionID, db.Data.Serial) @@ -478,7 +478,7 @@ func (db *DeploymentState) IsDeploymentMetadataService() bool { } func (db *DeploymentState) isDeploymentMetadataService() bool { - return db.mode == ReadingDMS || db.mode == WritingDMS + return db.mode == WritingDMS } // RecordsHistory reports whether this mode records resource operations with DMS. @@ -486,7 +486,7 @@ func (db *DeploymentState) RecordsHistory() bool { db.mu.Lock() defer db.mu.Unlock() switch db.mode { - case ReadingDMS, WritingDMS, WritingFileWithHistory, WritingFileWithHistoryEnforced: + case WritingDMS, WritingFileWithHistory, WritingFileWithHistoryEnforced: return true default: return false @@ -519,7 +519,6 @@ type WithRecovery bool // Open loads deployment state and enables reads or writes according to mode: // - ReadingFile: reads resources.json; state writes are disabled. // - WritingFile: reads resources.json and opens a WAL for state writes. -// - ReadingDMS: reads resources from DMS; state writes are disabled. // - WritingDMS: reads and writes resources through DMS; no WAL is opened. // - WritingFileWithHistory: behaves like WritingFile and also records DMS history. // History errors do not fail the deployment. @@ -600,7 +599,7 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode M db.stateIDs[key] = entry.ID } - recording := mode == ReadingDMS || mode == WritingDMS + recording := mode == WritingDMS walPath := db.Path + walSuffix _, err = os.Stat(walPath) @@ -686,8 +685,8 @@ Run "databricks bundle destroy" first, then deploy again with deployment history db.VersionID = dmsDeployment.LastVersionID } - if mode.isWritable() { - return db.initializeWrites(mode) + if mode.isWritable() && mode != WritingDMS { + return db.initializeFileWrites(mode) } return nil @@ -861,6 +860,9 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported db.walFile = nil err = db.replayWAL(ctx) } else if db.mode == WritingDMS { + db.Data.Lineage = db.GetOrInitLineage() + db.Data.StateVersion = currentStateVersion + db.Data.CLIVersion = build.GetInfo().Version // replayWAL is what normally persists the file. Without one, write the tombstone here so // the header (lineage, state version, CLI version) still lands on disk - even when no // operations were recorded, since the deployment and any version this run created exist @@ -894,52 +896,45 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported return state, err } -// UpgradeToWrite selects a write policy without re-reading resource state. +// UpgradeToWrite opens a WAL for file writes without re-reading resource state. +// DMS state is writable from Open, so upgrading it is a no-op. func (db *DeploymentState) UpgradeToWrite(mode Mode) error { db.mu.Lock() defer db.mu.Unlock() if db.Path == "" { return errors.New("internal error: DeploymentState must be opened first") } + if db.mode == WritingDMS { + return nil + } if db.mode.isWritable() { return errors.New("internal error: DeploymentState is already open for write") } - if !mode.isWritable() { + if !mode.isWritable() || mode == WritingDMS { return fmt.Errorf("invalid deployment state write mode: %d", mode) } - // Keep the state source unchanged: DMS reads require DMS writes; file reads require file writes. - if (db.mode == ReadingDMS) != (mode == WritingDMS) { - return errors.New("cannot change the deployment state source when upgrading to write") - } - if err := db.initializeWrites(mode); err != nil { + if err := db.initializeFileWrites(mode); err != nil { return err } db.mode = mode return nil } -// initializeWrites validates write compatibility and prepares DMS metadata or a WAL. +// initializeFileWrites validates file write compatibility and opens a WAL. // Callers must hold db.mu. -func (db *DeploymentState) initializeWrites(mode Mode) error { +func (db *DeploymentState) initializeFileWrites(mode Mode) error { if err := assertNoUnsupportedFeatures(db.Data.Features); err != nil { return err } _, recorded := db.Data.Features[FeatureDeploymentHistory] _, bestEffort := db.Data.Features[FeatureDeploymentHistoryBestEffort] - if recorded != (mode == WritingDMS) { + if recorded { return errors.New("cannot change the deployment state source when upgrading to write") } if bestEffort && mode != WritingFileWithHistory && mode != WritingFileWithHistoryEnforced { return errors.New("deployment_history_best_effort state requires writing with deployment history enabled") } - if mode == WritingDMS { - db.Data.Lineage = db.GetOrInitLineage() - db.Data.StateVersion = currentStateVersion - db.Data.CLIVersion = build.GetInfo().Version - return nil - } - walPath := db.Path + walSuffix if err := os.MkdirAll(filepath.Dir(walPath), 0o755); err != nil { return fmt.Errorf("failed to create state directory: %w", err) @@ -997,7 +992,7 @@ func (db *DeploymentState) AssertOpenedForReadOrWrite() { func (db *DeploymentState) AssertOpenedForRead() { db.AssertOpenedForReadOrWrite() - if db.mode.isWritable() { + if db.mode.isWritable() && db.mode != WritingDMS { panic("internal error: DeploymentState must be opened in read mode") } } @@ -1084,7 +1079,7 @@ func (db *DeploymentState) unlockedSave() error { // Data to persist in the remote resources.json file. func (db *DeploymentState) dataForFile() Database { - if db.mode == ReadingDMS || db.mode == WritingDMS { + if db.mode == WritingDMS { header := db.Data.Header if header.Features == nil { header.Features = make(map[string]struct{}) diff --git a/bundle/direct/dstate/state_test.go b/bundle/direct/dstate/state_test.go index 6bf51304077..2099ef7adbd 100644 --- a/bundle/direct/dstate/state_test.go +++ b/bundle/direct/dstate/state_test.go @@ -69,7 +69,7 @@ func TestDataForFileSerial(t *testing.T) { VersionID: tt.versionID, } if tt.recording { - db.mode = ReadingDMS + db.mode = WritingDMS } db.Data.State["resources.jobs.foo"] = ResourceEntry{ID: "123", State: json.RawMessage(`{"name":"foo"}`)} @@ -183,7 +183,7 @@ func TestOpenUsesDMSVersionForSerial(t *testing.T) { ctx := cmdctx.SetWorkspaceClient(t.Context(), client) var db DeploymentState - err = db.Open(ctx, path, ReadingDMS, WithRecovery(false), OpenDmsArgs{ + err = db.Open(ctx, path, WritingDMS, WithRecovery(false), OpenDmsArgs{ DeploymentID: tt.deploymentID, LastVersionID: tt.lastVersionID, }) if tt.wantErr != "" { diff --git a/bundle/phases/deploy.go b/bundle/phases/deploy.go index b906fa70bfe..3faed15ef4a 100644 --- a/bundle/phases/deploy.go +++ b/bundle/phases/deploy.go @@ -299,11 +299,7 @@ func Deploy(ctx context.Context, b *bundle.Bundle, outputHandler sync.OutputHand // Upgrade from read (opened by process.go, or by Migrate) to write mode. After approval // and the migration commit above, so a declined deploy never opens a WAL it must discard, // and the migration's tf+1 push lands before this advances the WAL header to tf+2. - writeMode := dstate.WritingFile - if b.DeploymentBundle.StateDB.IsDeploymentMetadataService() { - writeMode = dstate.WritingDMS - } - if err := b.DeploymentBundle.StateDB.UpgradeToWrite(writeMode); err != nil { + if err := b.DeploymentBundle.StateDB.UpgradeToWrite(dstate.WritingFile); err != nil { logdiag.LogError(ctx, err) return } diff --git a/bundle/phases/destroy.go b/bundle/phases/destroy.go index 6d4a1e6e049..7055fc3cadf 100644 --- a/bundle/phases/destroy.go +++ b/bundle/phases/destroy.go @@ -348,11 +348,7 @@ func Destroy(ctx context.Context, b *bundle.Bundle) { } // Upgrade from read (opened by process.go, or by Migrate) to write mode - writeMode := dstate.WritingFile - if b.DeploymentBundle.StateDB.IsDeploymentMetadataService() { - writeMode = dstate.WritingDMS - } - if err := b.DeploymentBundle.StateDB.UpgradeToWrite(writeMode); err != nil { + if err := b.DeploymentBundle.StateDB.UpgradeToWrite(dstate.WritingFile); err != nil { logdiag.LogError(ctx, err) return } diff --git a/cmd/bundle/utils/process.go b/cmd/bundle/utils/process.go index d542652b05e..8fa68f32f5e 100644 --- a/cmd/bundle/utils/process.go +++ b/cmd/bundle/utils/process.go @@ -406,7 +406,7 @@ func ProcessBundleRet(cmd *cobra.Command, opts ProcessOptions) (b *bundle.Bundle if !cmdctx.HasWorkspaceClient(ctx) { ctx = cmdctx.SetWorkspaceClient(ctx, b.WorkspaceClient(ctx)) } - if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.ReadingDMS, dstate.WithRecovery(false), dstate.OpenDmsArgs{DeploymentID: dmsDeploymentID, LastVersionID: lastVersionID}); err != nil { + if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WritingDMS, dstate.WithRecovery(false), dstate.OpenDmsArgs{DeploymentID: dmsDeploymentID, LastVersionID: lastVersionID}); err != nil { logdiag.LogError(ctx, err) return b, stateDesc, root.ErrAlreadyPrinted } @@ -657,7 +657,7 @@ func OpenDirectStateForRead(ctx context.Context, b *bundle.Bundle, stateDesc *st if !cmdctx.HasWorkspaceClient(ctx) { ctx = cmdctx.SetWorkspaceClient(ctx, b.WorkspaceClient(ctx)) } - if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.ReadingDMS, dstate.WithRecovery(false), dstate.OpenDmsArgs{DeploymentID: dmsDeploymentID, LastVersionID: lastVersionID}); err != nil { + if err := b.DeploymentBundle.StateDB.Open(ctx, localPath, dstate.WritingDMS, dstate.WithRecovery(false), dstate.OpenDmsArgs{DeploymentID: dmsDeploymentID, LastVersionID: lastVersionID}); err != nil { return err } return enforceDeploymentHistorySetting(ctx, b, stateDesc, false) From 1f230f82afba732708ad8e5440900427f2dc2622 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:22:50 +0000 Subject: [PATCH 25/52] Name WAL setup explicitly and spell out mode branches --- bundle/direct/dstate/state.go | 37 ++++++++++++++++++++++------------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 4cb8a626646..9b5e556d201 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -685,10 +685,14 @@ Run "databricks bundle destroy" first, then deploy again with deployment history db.VersionID = dmsDeployment.LastVersionID } - if mode.isWritable() && mode != WritingDMS { - return db.initializeFileWrites(mode) + switch mode { + case ReadingFile: + return nil + case WritingDMS: + return nil + case WritingFile, WritingFileWithHistory, WritingFileWithHistoryEnforced: + return db.setupWal(mode) } - return nil } @@ -904,25 +908,30 @@ func (db *DeploymentState) UpgradeToWrite(mode Mode) error { if db.Path == "" { return errors.New("internal error: DeploymentState must be opened first") } - if db.mode == WritingDMS { + switch db.mode { + case ReadingFile: + // Set up the requested file write mode below. + case WritingDMS: return nil - } - if db.mode.isWritable() { + case WritingFile, WritingFileWithHistory, WritingFileWithHistoryEnforced: return errors.New("internal error: DeploymentState is already open for write") } - if !mode.isWritable() || mode == WritingDMS { + + switch mode { + case WritingFile, WritingFileWithHistory, WritingFileWithHistoryEnforced: + if err := db.setupWal(mode); err != nil { + return err + } + db.mode = mode + return nil + default: return fmt.Errorf("invalid deployment state write mode: %d", mode) } - if err := db.initializeFileWrites(mode); err != nil { - return err - } - db.mode = mode - return nil } -// initializeFileWrites validates file write compatibility and opens a WAL. +// setupWal validates file write compatibility and opens a WAL. // Callers must hold db.mu. -func (db *DeploymentState) initializeFileWrites(mode Mode) error { +func (db *DeploymentState) setupWal(mode Mode) error { if err := assertNoUnsupportedFeatures(db.Data.Features); err != nil { return err } From cc1697356476d0dd611e846b73225e92d116d913 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:33:08 +0000 Subject: [PATCH 26/52] Preserve the ordinary WAL path and separate history branches --- bundle/direct/dstate/state.go | 123 +++++++++++++++++++--------------- 1 file changed, 69 insertions(+), 54 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 9b5e556d201..0c84cbe8cca 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -32,13 +32,14 @@ import ( const ( // currentStateVersion is the schema version written by this CLI and the version // legacy states are migrated up to on load. - currentStateVersion = 3 - initialBufferSize = 64 * 1024 - maxWalEntrySize = 10 * 1024 * 1024 - walSuffix = ".wal" - historyDrainTimeoutPerOperation = 5 * time.Second + currentStateVersion = 3 + initialBufferSize = 64 * 1024 + maxWalEntrySize = 10 * 1024 * 1024 + walSuffix = ".wal" ) +const historyDrainTimeoutPerOperation = 5 * time.Second + // FeatureDeploymentHistory marks a state whose resources are also recorded with the // deployment metadata service. Both stores are kept in step, so the marker is what tells a // reader the two already agree. A CLI that does not know the name refuses the state rather @@ -328,6 +329,8 @@ func (db *DeploymentState) SaveState(ctx context.Context, key, newID string, sta } switch db.mode { + case WritingFile: + // Use the normal WAL path below. case WritingDMS: db.Data.State[key] = entry db.stateIDs[key] = newID @@ -338,12 +341,7 @@ func (db *DeploymentState) SaveState(ctx context.Context, key, newID string, sta } buf.RecordOperation(ctx, key, false, newID, recorded) } - - case WritingFile: - if err := appendJSONLine(db.walFile, WALEntry{Key: key, Value: &entry}); err != nil { - return err - } - db.stateIDs[key] = newID + return nil case WritingFileWithHistory, WritingFileWithHistoryEnforced: // File-backed modes record history only after the WAL append succeeds. @@ -358,11 +356,17 @@ func (db *DeploymentState) SaveState(ctx context.Context, key, newID string, sta } buf.RecordOperation(ctx, key, false, newID, recorded) } + return nil default: return fmt.Errorf("cannot save state in mode %d", db.mode) } - return nil + + err = appendJSONLine(db.walFile, WALEntry{Key: key, Value: &entry}) + if err == nil { + db.stateIDs[key] = newID + } + return err } // DeleteState drops the resource's state entry: the resource is gone. inProgress records the @@ -380,18 +384,15 @@ func (db *DeploymentState) DeleteState(ctx context.Context, key string, inProgre deletedID := db.stateIDs[key] switch db.mode { + case WritingFile: + // Use the normal WAL path below. case WritingDMS: delete(db.Data.State, key) if buf := db.operationBuffer; buf != nil { buf.RecordOperation(ctx, key, inProgress, deletedID, nil) } delete(db.stateIDs, key) - - case WritingFile: - if err := appendJSONLine(db.walFile, WALEntry{Key: key}); err != nil { - return err - } - delete(db.stateIDs, key) + return nil case WritingFileWithHistory, WritingFileWithHistoryEnforced: if err := appendJSONLine(db.walFile, WALEntry{Key: key}); err != nil { @@ -401,10 +402,17 @@ func (db *DeploymentState) DeleteState(ctx context.Context, key string, inProgre buf.RecordOperation(ctx, key, inProgress, deletedID, nil) } delete(db.stateIDs, key) + return nil default: return fmt.Errorf("cannot delete state in mode %d", db.mode) } + + if err := appendJSONLine(db.walFile, WALEntry{Key: key}); err != nil { + return err + } + + delete(db.stateIDs, key) return nil } @@ -468,6 +476,11 @@ func (db *DeploymentState) GetSerial() int { } } +// isDeploymentMetadataService is IsDeploymentMetadataService for callers already holding db.mu. +func (db *DeploymentState) isDeploymentMetadataService() bool { + return db.mode == WritingDMS +} + // IsDeploymentMetadataService reports whether this deployment's resource state lives in the // deployment metadata service rather than the state file. The state's features are the source of // truth; valid only after Open. @@ -477,10 +490,6 @@ func (db *DeploymentState) IsDeploymentMetadataService() bool { return db.isDeploymentMetadataService() } -func (db *DeploymentState) isDeploymentMetadataService() bool { - return db.mode == WritingDMS -} - // RecordsHistory reports whether this mode records resource operations with DMS. func (db *DeploymentState) RecordsHistory() bool { db.mu.Lock() @@ -543,7 +552,10 @@ func (db *DeploymentState) Open(ctx context.Context, path string, mode Mode, wit err := db.unlockedOpen(ctx, path, mode, withRecovery, dmsDeployment) if err != nil { - // A failed open must leave the receiver closed so it can be retried. + // A failed open must leave the receiver closed. unlockedOpen assigns + // db.Path before every fallible step, so without this the receiver stays + // half-initialized and the next Open on it hits the panic above instead + // of reporting the real error. db.reset() } return err @@ -818,14 +830,16 @@ func (db *DeploymentState) mergeWalIntoState(ctx context.Context) (bool, error) } } - // Persist resource entries after the header, or feature changes in the header itself. - // Enforced history can add its feature marker without any resource writes; that marker - // still needs to reach resources.json before the WAL is removed. - hasEntries := lineNumber > 1 || lineNumber == 1 && !maps.Equal(db.Data.Features, newFeatures) + hasEntries := lineNumber > 1 + // Enforced history can change the header's features without writing resource entries. + // Persist that marker too, before removing the WAL. + if lineNumber == 1 && !maps.Equal(db.Data.Features, newFeatures) { + hasEntries = true + } - // Only advance the serial when the WAL carried changes, because the caller + // Only advance the serial when the WAL carried entries, because the caller // (replayWAL) persists the new state file only in that case. A header-only - // WAL with unchanged features is a deploy that started but committed nothing; advancing the serial + // WAL is a deploy that started but committed nothing; advancing the serial // for it leaves the in-memory serial ahead of the persisted one, so the // next deploy writes its WAL header at serial+2 and recovery rejects it as // "ahead of expected". See acceptance/bundle/deploy/wal/header-only-wal. @@ -874,22 +888,18 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported err = db.unlockedSave() } - // File state is persisted before waiting on optional history recording. - if db.operationBuffer != nil { - bestEffort := db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced - drainCtx := ctx - cancel := func() {} - if bestEffort { - drainCtx, cancel = context.WithTimeout(ctx, db.historyDrainTimeout) + // Wait until all operations are recorded in the service. + if db.isDeploymentMetadataService() && db.operationBuffer != nil { + if drainErr := db.operationBuffer.Drain(ctx); drainErr != nil { + err = errors.Join(err, drainErr) } + } else if (db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced) && db.operationBuffer != nil { + // File state is persisted before waiting on optional history recording. + drainCtx, cancel := context.WithTimeout(ctx, db.historyDrainTimeout) drainErr := db.operationBuffer.Drain(drainCtx) cancel() if drainErr != nil { - if bestEffort { - log.Warnf(ctx, "Failed to record deployment history: %s", drainErr) - } else { - err = errors.Join(err, drainErr) - } + log.Warnf(ctx, "Failed to record deployment history: %s", drainErr) } } @@ -952,27 +962,32 @@ func (db *DeploymentState) setupWal(mode Mode) error { if err != nil { return fmt.Errorf("failed to open WAL file %s: %w", walPath, err) } - features := maps.Clone(db.Data.Features) - if mode == WritingFileWithHistoryEnforced { - if features == nil { - features = make(map[string]struct{}) - } - features[FeatureDeploymentHistoryBestEffort] = struct{}{} - } + db.walFile = walFile walHead := Header{ Lineage: db.GetOrInitLineage(), Serial: db.Data.Serial + 1, StateVersion: currentStateVersion, CLIVersion: build.GetInfo().Version, - Features: features, } - if err := appendJSONLine(walFile, walHead); err != nil { - // O_EXCL guarantees this is ours; an incomplete header cannot be recovered. - walFile.Close() + if mode == WritingFileWithHistory || mode == WritingFileWithHistoryEnforced { + walHead.Features = maps.Clone(db.Data.Features) + if mode == WritingFileWithHistoryEnforced { + if walHead.Features == nil { + walHead.Features = make(map[string]struct{}) + } + walHead.Features[FeatureDeploymentHistoryBestEffort] = struct{}{} + } + } + if err := appendJSONLine(db.walFile, walHead); err != nil { + // Remove the WAL created just above: O_EXCL means it is ours, and + // without a header it carries nothing to recover. Leaving it makes + // every later open fail, with recovery on the header parse and + // without it on the unexpected WAL file. + db.walFile.Close() + db.walFile = nil os.Remove(walPath) return fmt.Errorf("failed to write WAL header to %s: %w", walPath, err) } - db.walFile = walFile return nil } @@ -1088,7 +1103,7 @@ func (db *DeploymentState) unlockedSave() error { // Data to persist in the remote resources.json file. func (db *DeploymentState) dataForFile() Database { - if db.mode == WritingDMS { + if db.isDeploymentMetadataService() { header := db.Data.Header if header.Features == nil { header.Features = make(map[string]struct{}) From efae1627e86241963a94ffa911f766df1c183920 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:37:33 +0000 Subject: [PATCH 27/52] Remove redundant operation buffer write assertion --- bundle/direct/dstate/state.go | 1 - 1 file changed, 1 deletion(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 0c84cbe8cca..6b14706cc53 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -183,7 +183,6 @@ type WALEntry struct { // (after approval) - which is why it is not an Open option. It also records the id a first deploy // just created, which Open could not know, so CompleteVersion later has it. func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploymentID string, versionID, operationCount int) { - db.AssertOpenedForWrite() db.mu.Lock() defer db.mu.Unlock() options := dms.BufferOptions{ From 05e55a92e23d4615d3d29977b2e57951e954e601 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:42:13 +0000 Subject: [PATCH 28/52] Reset state unconditionally when completing a version --- bundle/direct/dstate/state.go | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 6b14706cc53..98f623aaf0e 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -236,9 +236,7 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b db.versionCompleted = true deploymentID, versionID := db.DeploymentID, db.VersionID bestEffort := db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced - if db.Path == "" { - db.reset() - } + db.reset() db.mu.Unlock() // A recording failure fails the version even when the caller counted the deploy a success: the From 76c034352a5d6e0c5d981418ae8d616cc12592f7 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:47:02 +0000 Subject: [PATCH 29/52] Inline normal file persistence in save and delete mode cases --- bundle/direct/dstate/state.go | 25 ++++++++++--------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 98f623aaf0e..3ee6cb32b99 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -327,7 +327,11 @@ func (db *DeploymentState) SaveState(ctx context.Context, key, newID string, sta switch db.mode { case WritingFile: - // Use the normal WAL path below. + err = appendJSONLine(db.walFile, WALEntry{Key: key, Value: &entry}) + if err == nil { + db.stateIDs[key] = newID + } + return err case WritingDMS: db.Data.State[key] = entry db.stateIDs[key] = newID @@ -358,12 +362,6 @@ func (db *DeploymentState) SaveState(ctx context.Context, key, newID string, sta default: return fmt.Errorf("cannot save state in mode %d", db.mode) } - - err = appendJSONLine(db.walFile, WALEntry{Key: key, Value: &entry}) - if err == nil { - db.stateIDs[key] = newID - } - return err } // DeleteState drops the resource's state entry: the resource is gone. inProgress records the @@ -382,7 +380,11 @@ func (db *DeploymentState) DeleteState(ctx context.Context, key string, inProgre switch db.mode { case WritingFile: - // Use the normal WAL path below. + if err := appendJSONLine(db.walFile, WALEntry{Key: key}); err != nil { + return err + } + delete(db.stateIDs, key) + return nil case WritingDMS: delete(db.Data.State, key) if buf := db.operationBuffer; buf != nil { @@ -404,13 +406,6 @@ func (db *DeploymentState) DeleteState(ctx context.Context, key string, inProgre default: return fmt.Errorf("cannot delete state in mode %d", db.mode) } - - if err := appendJSONLine(db.walFile, WALEntry{Key: key}); err != nil { - return err - } - - delete(db.stateIDs, key) - return nil } func (db *DeploymentState) GetResourceEntry(key string) (ResourceEntry, bool) { From 198e51da9ad7098cf83551bc1903ba78ef0d5fb7 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:48:50 +0000 Subject: [PATCH 30/52] Clarify which recording fields survive finalization and why --- bundle/direct/dstate/state.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 3ee6cb32b99..10be39de299 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -563,7 +563,8 @@ func (db *DeploymentState) reset() { db.Data = Database{} db.stateIDs = nil db.historyDrainTimeout = 0 - // Retain recording state only while version completion is pending. + // Keep the buffer, mode, and IDs after Finalize so CompleteVersion can use them. + // Clear them if no version was started, or CompleteVersion has copied what it needs. if db.operationBuffer == nil || db.versionCompleted { db.mode = ReadingFile db.operationBuffer = nil From d8206d6edac0e7e5990775740bc8c5ca80455841 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:50:59 +0000 Subject: [PATCH 31/52] Remove optional best-effort marker check when opening state --- bundle/direct/dstate/mode_test.go | 3 +-- bundle/direct/dstate/state.go | 4 ---- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index b473891234d..36d5f3537db 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -428,8 +428,7 @@ func TestOpenRejectsIncompatibleModesAndFeatures(t *testing.T) { wantErr string }{ {"unknown feature", dstate.ReadingFile, `{"future":{}}`, "requires features"}, - {"conflicting features", dstate.ReadingFile, `{"deployment_history":{},"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, - {"DMS with best effort history", dstate.WritingDMS, `{"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, + {"file reading DMS with best effort marker", dstate.ReadingFile, `{"deployment_history":{},"deployment_history_best_effort":{}}`, "not supported"}, {"file writing DMS", dstate.WritingFile, `{"deployment_history":{}}`, "not supported"}, {"file writing best effort history", dstate.WritingFile, `{"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, } { diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 10be39de299..d9db7fbe159 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -637,10 +637,6 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode M if recording && recorded && db.Data.Serial > dmsDeployment.LastVersionID { return fmt.Errorf("inconsistent deployment state: resource state serial (%d) is ahead of deployment history version (%d)", db.Data.Serial, dmsDeployment.LastVersionID) } - _, bestEffort := db.Data.Features[FeatureDeploymentHistoryBestEffort] - if bestEffort && (recorded || recording) { - return errors.New("deployment_history_best_effort file state cannot be opened as DMS-authoritative state") - } switch { case recording && !recorded && len(db.Data.State) == 0: if db.Data.Features == nil { From 33449b71e2b60e64b04feaa7b0591d46e2408e81 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:57:05 +0000 Subject: [PATCH 32/52] Use stored modes for write setup and explicit drain policies --- bundle/direct/dstate/mode_test.go | 40 +++++++-------- bundle/direct/dstate/state.go | 74 ++++++++++++++-------------- bundle/direct/dstate/state_test.go | 2 +- bundle/migrate/build_state_test.go | 2 +- bundle/phases/deploy.go | 3 +- bundle/phases/destroy.go | 3 +- bundle/statemgmt/direct_migration.go | 2 +- cmd/bundle/deployment/migrate.go | 2 +- 8 files changed, 63 insertions(+), 65 deletions(-) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index 36d5f3537db..3c29a205800 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -200,7 +200,7 @@ func TestDMSReadsAndWritesWithoutWAL(t *testing.T) { require.NoError(t, db.Open(ctx, path, dstate.WritingDMS, false, dstate.OpenDmsArgs{DeploymentID: "dep-1", LastVersionID: 7})) assert.Equal(t, "remote-id", db.GetResourceID("resources.jobs.remote")) assert.Equal(t, 7, db.GetSerial()) - require.NoError(t, db.UpgradeToWrite(dstate.WritingDMS)) + require.NoError(t, db.UpgradeToWrite()) assert.NoFileExists(t, path+".wal") db.InitializeOperationBuffer(ctx, "dep-1", 8, 2) require.NoError(t, db.SaveState(ctx, "resources.jobs.new", "new-id", map[string]string{"name": "new"}, nil)) @@ -219,24 +219,20 @@ func TestDMSReadsAndWritesWithoutWAL(t *testing.T) { assert.Contains(t, marker.Features, dstate.FeatureDeploymentHistory) } -func TestUpgradeToWriteValidatesMode(t *testing.T) { +func TestUpgradeToWrite(t *testing.T) { for _, tc := range []struct { name string - read dstate.Mode - write dstate.Mode + mode dstate.Mode bestEffort bool wantErr string }{ - {"file", dstate.ReadingFile, dstate.WritingFile, false, ""}, - {"history", dstate.ReadingFile, dstate.WritingFileWithHistory, false, ""}, - {"enforced history", dstate.ReadingFile, dstate.WritingFileWithHistoryEnforced, false, ""}, - {"DMS", dstate.WritingDMS, dstate.WritingDMS, false, ""}, - {"file to DMS", dstate.ReadingFile, dstate.WritingDMS, false, "invalid"}, - {"DMS needs no WAL upgrade", dstate.WritingDMS, dstate.WritingFile, false, ""}, - {"reading mode", dstate.ReadingFile, dstate.ReadingFile, false, "invalid"}, - {"unknown mode", dstate.ReadingFile, dstate.Mode(99), false, "invalid"}, - {"enforcement retained", dstate.ReadingFile, dstate.WritingFileWithHistory, true, ""}, - {"enforcement bypass", dstate.ReadingFile, dstate.WritingFile, true, "deployment_history_best_effort"}, + {"file", dstate.ReadingFile, false, ""}, + {"DMS", dstate.WritingDMS, false, ""}, + {"history", dstate.WritingFileWithHistory, false, ""}, + {"enforced history", dstate.WritingFileWithHistoryEnforced, false, ""}, + {"file already writable", dstate.WritingFile, false, "already open for write"}, + {"enforcement retained", dstate.WritingFileWithHistory, true, ""}, + {"enforcement bypass", dstate.ReadingFile, true, "deployment_history_best_effort"}, } { t.Run(tc.name, func(t *testing.T) { ctx := modeContext(t, nil) @@ -245,12 +241,14 @@ func TestUpgradeToWriteValidatesMode(t *testing.T) { require.NoError(t, os.WriteFile(path, []byte(`{"state_version":3,"features":{"deployment_history_best_effort":{}},"state":{}}`), 0o600)) } var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, tc.read, false, dstate.OpenDmsArgs{})) - err := db.UpgradeToWrite(tc.write) + require.NoError(t, db.Open(ctx, path, tc.mode, false, dstate.OpenDmsArgs{})) + err := db.UpgradeToWrite() if tc.wantErr != "" { require.ErrorContains(t, err, tc.wantErr) - assert.NotPanics(t, db.AssertOpenedForRead) - assert.NoFileExists(t, path+".wal") + if tc.mode == dstate.ReadingFile { + assert.NotPanics(t, db.AssertOpenedForRead) + assert.NoFileExists(t, path+".wal") + } } else { require.NoError(t, err) assert.NotPanics(t, db.AssertOpenedForWrite) @@ -410,7 +408,7 @@ func TestEnforcedHistorySurvivesWALRecovery(t *testing.T) { assert.Equal(t, 1, recovered.GetSerial()) assert.FileExists(t, recoveredPath) assert.NoFileExists(t, recoveredPath+".wal") - require.ErrorContains(t, recovered.UpgradeToWrite(dstate.WritingFile), "deployment_history_best_effort") + require.ErrorContains(t, recovered.UpgradeToWrite(), "deployment_history_best_effort") if withResource { assert.Equal(t, "id-1", recovered.GetResourceID("resources.jobs.job")) } @@ -491,10 +489,10 @@ func TestFailedUpgradeLeavesReadModeAvailable(t *testing.T) { var db dstate.DeploymentState require.NoError(t, db.Open(t.Context(), path, dstate.ReadingFile, true, dstate.OpenDmsArgs{})) require.NoError(t, os.WriteFile(path+".wal", []byte("occupied"), 0o600)) - require.Error(t, db.UpgradeToWrite(dstate.WritingFile)) + require.Error(t, db.UpgradeToWrite()) assert.NotPanics(t, db.AssertOpenedForRead) require.NoError(t, os.Remove(path+".wal")) - require.NoError(t, db.UpgradeToWrite(dstate.WritingFile)) + require.NoError(t, db.UpgradeToWrite()) _, err := db.Finalize(t.Context()) require.NoError(t, err) } diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index d9db7fbe159..10af69a05ea 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -431,10 +431,8 @@ func (db *DeploymentState) GetResourceID(key string) string { return db.stateIDs[key] } -// StateCLIVersion returns the CLI version that last wrote the state, or an empty -// string if the state does not record one (a fresh state that this CLI has not -// written yet). It is the version stored in the on-disk header, not the running -// build's version. +// StateCLIVersion returns the CLI version in the current state header. +// File state retains the last writer's version; DMS stamps the running version during Open. func (db *DeploymentState) StateCLIVersion() string { db.AssertOpenedForReadOrWrite() db.mu.Lock() @@ -690,9 +688,11 @@ Run "databricks bundle destroy" first, then deploy again with deployment history case ReadingFile: return nil case WritingDMS: + db.Data.StateVersion = currentStateVersion + db.Data.CLIVersion = build.GetInfo().Version return nil case WritingFile, WritingFileWithHistory, WritingFileWithHistoryEnforced: - return db.setupWal(mode) + return db.setupWal() } return nil } @@ -867,9 +867,8 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported db.walFile = nil err = db.replayWAL(ctx) } else if db.mode == WritingDMS { + // Create a lineage when persisting it so saved plans use a stable value. db.Data.Lineage = db.GetOrInitLineage() - db.Data.StateVersion = currentStateVersion - db.Data.CLIVersion = build.GetInfo().Version // replayWAL is what normally persists the file. Without one, write the tombstone here so // the header (lineage, state version, CLI version) still lands on disk - even when no // operations were recorded, since the deployment and any version this run created exist @@ -878,17 +877,24 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported } // Wait until all operations are recorded in the service. - if db.isDeploymentMetadataService() && db.operationBuffer != nil { - if drainErr := db.operationBuffer.Drain(ctx); drainErr != nil { - err = errors.Join(err, drainErr) + switch db.mode { + case ReadingFile, WritingFile: + // File-only modes have no history to drain. + case WritingDMS: + if db.operationBuffer != nil { + if drainErr := db.operationBuffer.Drain(ctx); drainErr != nil { + err = errors.Join(err, drainErr) + } } - } else if (db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced) && db.operationBuffer != nil { - // File state is persisted before waiting on optional history recording. - drainCtx, cancel := context.WithTimeout(ctx, db.historyDrainTimeout) - drainErr := db.operationBuffer.Drain(drainCtx) - cancel() - if drainErr != nil { - log.Warnf(ctx, "Failed to record deployment history: %s", drainErr) + case WritingFileWithHistory, WritingFileWithHistoryEnforced: + if db.operationBuffer != nil { + // File state is persisted before waiting on optional history recording. + drainCtx, cancel := context.WithTimeout(ctx, db.historyDrainTimeout) + drainErr := db.operationBuffer.Drain(drainCtx) + cancel() + if drainErr != nil { + log.Warnf(ctx, "Failed to record deployment history: %s", drainErr) + } } } @@ -899,38 +905,34 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported return state, err } -// UpgradeToWrite opens a WAL for file writes without re-reading resource state. -// DMS state is writable from Open, so upgrading it is a no-op. -func (db *DeploymentState) UpgradeToWrite(mode Mode) error { +// UpgradeToWrite opens a WAL for file state without re-reading resources. +// DMS and file-history modes are already writable from Open, so upgrading them is a no-op. +func (db *DeploymentState) UpgradeToWrite() error { db.mu.Lock() defer db.mu.Unlock() + if db.Path == "" { return errors.New("internal error: DeploymentState must be opened first") } switch db.mode { - case ReadingFile: - // Set up the requested file write mode below. - case WritingDMS: + case WritingDMS, WritingFileWithHistory, WritingFileWithHistoryEnforced: return nil - case WritingFile, WritingFileWithHistory, WritingFileWithHistoryEnforced: + case WritingFile: return errors.New("internal error: DeploymentState is already open for write") - } - - switch mode { - case WritingFile, WritingFileWithHistory, WritingFileWithHistoryEnforced: - if err := db.setupWal(mode); err != nil { + case ReadingFile: + db.mode = WritingFile + if err := db.setupWal(); err != nil { + db.mode = ReadingFile return err } - db.mode = mode return nil - default: - return fmt.Errorf("invalid deployment state write mode: %d", mode) } + return fmt.Errorf("invalid deployment state mode: %d", db.mode) } // setupWal validates file write compatibility and opens a WAL. // Callers must hold db.mu. -func (db *DeploymentState) setupWal(mode Mode) error { +func (db *DeploymentState) setupWal() error { if err := assertNoUnsupportedFeatures(db.Data.Features); err != nil { return err } @@ -939,7 +941,7 @@ func (db *DeploymentState) setupWal(mode Mode) error { if recorded { return errors.New("cannot change the deployment state source when upgrading to write") } - if bestEffort && mode != WritingFileWithHistory && mode != WritingFileWithHistoryEnforced { + if bestEffort && db.mode != WritingFileWithHistory && db.mode != WritingFileWithHistoryEnforced { return errors.New("deployment_history_best_effort state requires writing with deployment history enabled") } @@ -958,9 +960,9 @@ func (db *DeploymentState) setupWal(mode Mode) error { StateVersion: currentStateVersion, CLIVersion: build.GetInfo().Version, } - if mode == WritingFileWithHistory || mode == WritingFileWithHistoryEnforced { + if db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced { walHead.Features = maps.Clone(db.Data.Features) - if mode == WritingFileWithHistoryEnforced { + if db.mode == WritingFileWithHistoryEnforced { if walHead.Features == nil { walHead.Features = make(map[string]struct{}) } diff --git a/bundle/direct/dstate/state_test.go b/bundle/direct/dstate/state_test.go index 2099ef7adbd..e2aa1be6e4d 100644 --- a/bundle/direct/dstate/state_test.go +++ b/bundle/direct/dstate/state_test.go @@ -390,7 +390,7 @@ func TestGetOrInitLineageReadableBeforeWriteAndPersisted(t *testing.T) { // Upgrading to write reuses the same lineage (it goes into the WAL header), // and a write makes it durable. - require.NoError(t, db.UpgradeToWrite(WritingFile)) + require.NoError(t, db.UpgradeToWrite()) require.NoError(t, db.SaveState(t.Context(), "jobs.my_job", "123", map[string]string{}, nil)) mustFinalize(t, &db) diff --git a/bundle/migrate/build_state_test.go b/bundle/migrate/build_state_test.go index f9c7cac9519..a495436008c 100644 --- a/bundle/migrate/build_state_test.go +++ b/bundle/migrate/build_state_test.go @@ -48,7 +48,7 @@ func runBuildStateFromTF( var db dstate.DeploymentState db.OpenWithData(statePath, dstate.NewDatabase("lineage", 1)) - require.NoError(t, db.UpgradeToWrite(dstate.WritingFile)) + require.NoError(t, db.UpgradeToWrite()) _, err = migrate.BuildStateFromTF(t.Context(), &root, adapters, &db, tfAttrs, tfIDs, "") require.NoError(t, err) diff --git a/bundle/phases/deploy.go b/bundle/phases/deploy.go index 3faed15ef4a..dca864ae670 100644 --- a/bundle/phases/deploy.go +++ b/bundle/phases/deploy.go @@ -15,7 +15,6 @@ import ( "github.com/databricks/cli/bundle/deploy/metadata" "github.com/databricks/cli/bundle/deploy/snapshot" "github.com/databricks/cli/bundle/deployplan" - "github.com/databricks/cli/bundle/direct/dstate" "github.com/databricks/cli/bundle/libraries" "github.com/databricks/cli/bundle/metrics" "github.com/databricks/cli/bundle/permissions" @@ -299,7 +298,7 @@ func Deploy(ctx context.Context, b *bundle.Bundle, outputHandler sync.OutputHand // Upgrade from read (opened by process.go, or by Migrate) to write mode. After approval // and the migration commit above, so a declined deploy never opens a WAL it must discard, // and the migration's tf+1 push lands before this advances the WAL header to tf+2. - if err := b.DeploymentBundle.StateDB.UpgradeToWrite(dstate.WritingFile); err != nil { + if err := b.DeploymentBundle.StateDB.UpgradeToWrite(); err != nil { logdiag.LogError(ctx, err) return } diff --git a/bundle/phases/destroy.go b/bundle/phases/destroy.go index 7055fc3cadf..408b4fd3191 100644 --- a/bundle/phases/destroy.go +++ b/bundle/phases/destroy.go @@ -14,7 +14,6 @@ import ( "github.com/databricks/cli/bundle/deploy/files" "github.com/databricks/cli/bundle/deploy/lock" "github.com/databricks/cli/bundle/deployplan" - "github.com/databricks/cli/bundle/direct/dstate" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/diag" "github.com/databricks/cli/libs/dms" @@ -348,7 +347,7 @@ func Destroy(ctx context.Context, b *bundle.Bundle) { } // Upgrade from read (opened by process.go, or by Migrate) to write mode - if err := b.DeploymentBundle.StateDB.UpgradeToWrite(dstate.WritingFile); err != nil { + if err := b.DeploymentBundle.StateDB.UpgradeToWrite(); err != nil { logdiag.LogError(ctx, err) return } diff --git a/bundle/statemgmt/direct_migration.go b/bundle/statemgmt/direct_migration.go index 7a06d75aca9..4fefd8504ac 100644 --- a/bundle/statemgmt/direct_migration.go +++ b/bundle/statemgmt/direct_migration.go @@ -287,7 +287,7 @@ func convertTFStateToDirect(ctx context.Context, b *bundle.Bundle, tfState *migr return tempStatePath, false, nil, err } - if err := stateDB.UpgradeToWrite(dstate.WritingFile); err != nil { + if err := stateDB.UpgradeToWrite(); err != nil { return tempStatePath, false, nil, fmt.Errorf("upgrading state for apply: %w", err) } diff --git a/cmd/bundle/deployment/migrate.go b/cmd/bundle/deployment/migrate.go index c9b82705d5e..f329ec7962c 100644 --- a/cmd/bundle/deployment/migrate.go +++ b/cmd/bundle/deployment/migrate.go @@ -149,7 +149,7 @@ New bundles use the direct engine automatically, no migration is needed.`) return err } - if err := stateDB.UpgradeToWrite(dstate.WritingFile); err != nil { + if err := stateDB.UpgradeToWrite(); err != nil { return fmt.Errorf("upgrading state for apply: %w", err) } From e695feb94e3b64947cf7b87634039bdd09171f7e Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 10:58:53 +0000 Subject: [PATCH 33/52] Preserve DMS header timing for plans and downgrade warnings --- bundle/direct/dstate/state.go | 13 ++++++++----- bundle/direct/dstate/state_test.go | 2 ++ 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 10af69a05ea..02c79571041 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -431,8 +431,10 @@ func (db *DeploymentState) GetResourceID(key string) string { return db.stateIDs[key] } -// StateCLIVersion returns the CLI version in the current state header. -// File state retains the last writer's version; DMS stamps the running version during Open. +// StateCLIVersion returns the CLI version that last wrote the state, or an empty +// string if the state does not record one (a fresh state that this CLI has not +// written yet). It is the version stored in the on-disk header, not the running +// build's version. func (db *DeploymentState) StateCLIVersion() string { db.AssertOpenedForReadOrWrite() db.mu.Lock() @@ -688,8 +690,6 @@ Run "databricks bundle destroy" first, then deploy again with deployment history case ReadingFile: return nil case WritingDMS: - db.Data.StateVersion = currentStateVersion - db.Data.CLIVersion = build.GetInfo().Version return nil case WritingFile, WritingFileWithHistory, WritingFileWithHistoryEnforced: return db.setupWal() @@ -867,8 +867,11 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported db.walFile = nil err = db.replayWAL(ctx) } else if db.mode == WritingDMS { - // Create a lineage when persisting it so saved plans use a stable value. + // Stamp the header when persisting it: saved plans need a stable lineage, and + // the last writer's CLI version must remain available for downgrade warnings. db.Data.Lineage = db.GetOrInitLineage() + db.Data.StateVersion = currentStateVersion + db.Data.CLIVersion = build.GetInfo().Version // replayWAL is what normally persists the file. Without one, write the tombstone here so // the header (lineage, state version, CLI version) still lands on disk - even when no // operations were recorded, since the deployment and any version this run created exist diff --git a/bundle/direct/dstate/state_test.go b/bundle/direct/dstate/state_test.go index e2aa1be6e4d..2395ee349a2 100644 --- a/bundle/direct/dstate/state_test.go +++ b/bundle/direct/dstate/state_test.go @@ -167,6 +167,7 @@ func TestOpenUsesDMSVersionForSerial(t *testing.T) { path := filepath.Join(t.TempDir(), "resources.json") seed := map[string]any{ "state_version": currentStateVersion, + "cli_version": "99.0.0", "lineage": "test-lineage", "features": map[string]struct{}{FeatureDeploymentHistory: {}}, "state": map[string]ResourceEntry{}, @@ -193,6 +194,7 @@ func TestOpenUsesDMSVersionForSerial(t *testing.T) { } require.NoError(t, err) assert.Equal(t, tt.lastVersionID, db.GetSerial()) + assert.Equal(t, "99.0.0", db.StateCLIVersion(), "preserve the last writer for downgrade warnings") mustFinalize(t, &db) }) } From 8c94c94d2d67c3d3155b558c63c1dd38aa73db53 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 11:02:21 +0000 Subject: [PATCH 34/52] Inline DMS mode checks when opening deployment state --- bundle/direct/dstate/state.go | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 02c79571041..0388381e286 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -604,8 +604,6 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode M db.stateIDs[key] = entry.ID } - recording := mode == WritingDMS - walPath := db.Path + walSuffix _, err = os.Stat(walPath) switch { @@ -615,7 +613,7 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode M return fmt.Errorf("failed to stat WAL file %s: %w", walPath, err) default: // WAL exists switch { - case recording: + case mode == WritingDMS: // A recorded deployment writes no WAL, so finding one means this state was written by a // deployment that did not record history. Refuse rather than discard it: the file is the // only record of that deploy's writes, and recording is refused below anyway. @@ -634,21 +632,21 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, mode M // the config can neither start recording one that was not (its resources would be created a // second time) nor stop recording one that is (the service still holds it). _, recorded := db.Data.Features[FeatureDeploymentHistory] - if recording && recorded && db.Data.Serial > dmsDeployment.LastVersionID { + if mode == WritingDMS && recorded && db.Data.Serial > dmsDeployment.LastVersionID { return fmt.Errorf("inconsistent deployment state: resource state serial (%d) is ahead of deployment history version (%d)", db.Data.Serial, dmsDeployment.LastVersionID) } switch { - case recording && !recorded && len(db.Data.State) == 0: + case mode == WritingDMS && !recorded && len(db.Data.State) == 0: if db.Data.Features == nil { db.Data.Features = make(map[string]struct{}, 1) } db.Data.Features[FeatureDeploymentHistory] = struct{}{} recorded = true - case recording && !recorded: + case mode == WritingDMS && !recorded: return errors.New(`enabling experimental.deployment_history for an existing deployment is not supported Run "databricks bundle destroy" first, then deploy again with deployment history enabled`) - case !recording && recorded: + case mode != WritingDMS && recorded: return ErrUnsettingRecording } From a735d8e8b9d5baf4e02cdf0a724b5cd7b2bce916 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 11:04:29 +0000 Subject: [PATCH 35/52] Remove optional WAL mode checks and requested failure test --- bundle/direct/dstate/mode_test.go | 17 ++++++++++------- bundle/direct/dstate/state.go | 22 ++++++---------------- bundle/direct/dstate/state_test.go | 22 ---------------------- 3 files changed, 16 insertions(+), 45 deletions(-) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index 3c29a205800..670adeb043a 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -232,7 +232,8 @@ func TestUpgradeToWrite(t *testing.T) { {"enforced history", dstate.WritingFileWithHistoryEnforced, false, ""}, {"file already writable", dstate.WritingFile, false, "already open for write"}, {"enforcement retained", dstate.WritingFileWithHistory, true, ""}, - {"enforcement bypass", dstate.ReadingFile, true, "deployment_history_best_effort"}, + {"file upgrade with history marker", dstate.ReadingFile, true, ""}, + {"file open with history marker", dstate.WritingFile, true, "already open for write"}, } { t.Run(tc.name, func(t *testing.T) { ctx := modeContext(t, nil) @@ -245,16 +246,19 @@ func TestUpgradeToWrite(t *testing.T) { err := db.UpgradeToWrite() if tc.wantErr != "" { require.ErrorContains(t, err, tc.wantErr) - if tc.mode == dstate.ReadingFile { - assert.NotPanics(t, db.AssertOpenedForRead) - assert.NoFileExists(t, path+".wal") - } } else { require.NoError(t, err) assert.NotPanics(t, db.AssertOpenedForWrite) } _, err = db.Finalize(ctx) require.NoError(t, err) + if tc.bestEffort { + data, err := os.ReadFile(path) + require.NoError(t, err) + var persisted dstate.Database + require.NoError(t, json.Unmarshal(data, &persisted)) + assert.Contains(t, persisted.Features, dstate.FeatureDeploymentHistoryBestEffort) + } }) } } @@ -408,7 +412,7 @@ func TestEnforcedHistorySurvivesWALRecovery(t *testing.T) { assert.Equal(t, 1, recovered.GetSerial()) assert.FileExists(t, recoveredPath) assert.NoFileExists(t, recoveredPath+".wal") - require.ErrorContains(t, recovered.UpgradeToWrite(), "deployment_history_best_effort") + require.NoError(t, recovered.UpgradeToWrite()) if withResource { assert.Equal(t, "id-1", recovered.GetResourceID("resources.jobs.job")) } @@ -428,7 +432,6 @@ func TestOpenRejectsIncompatibleModesAndFeatures(t *testing.T) { {"unknown feature", dstate.ReadingFile, `{"future":{}}`, "requires features"}, {"file reading DMS with best effort marker", dstate.ReadingFile, `{"deployment_history":{},"deployment_history_best_effort":{}}`, "not supported"}, {"file writing DMS", dstate.WritingFile, `{"deployment_history":{}}`, "not supported"}, - {"file writing best effort history", dstate.WritingFile, `{"deployment_history_best_effort":{}}`, "deployment_history_best_effort"}, } { t.Run(tc.name, func(t *testing.T) { path := filepath.Join(t.TempDir(), "resources.json") diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 0388381e286..759d3f40eb4 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -931,20 +931,12 @@ func (db *DeploymentState) UpgradeToWrite() error { return fmt.Errorf("invalid deployment state mode: %d", db.mode) } -// setupWal validates file write compatibility and opens a WAL. +// setupWal opens a WAL and writes its header. // Callers must hold db.mu. func (db *DeploymentState) setupWal() error { if err := assertNoUnsupportedFeatures(db.Data.Features); err != nil { return err } - _, recorded := db.Data.Features[FeatureDeploymentHistory] - _, bestEffort := db.Data.Features[FeatureDeploymentHistoryBestEffort] - if recorded { - return errors.New("cannot change the deployment state source when upgrading to write") - } - if bestEffort && db.mode != WritingFileWithHistory && db.mode != WritingFileWithHistoryEnforced { - return errors.New("deployment_history_best_effort state requires writing with deployment history enabled") - } walPath := db.Path + walSuffix if err := os.MkdirAll(filepath.Dir(walPath), 0o755); err != nil { @@ -960,15 +952,13 @@ func (db *DeploymentState) setupWal() error { Serial: db.Data.Serial + 1, StateVersion: currentStateVersion, CLIVersion: build.GetInfo().Version, + Features: maps.Clone(db.Data.Features), } - if db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced { - walHead.Features = maps.Clone(db.Data.Features) - if db.mode == WritingFileWithHistoryEnforced { - if walHead.Features == nil { - walHead.Features = make(map[string]struct{}) - } - walHead.Features[FeatureDeploymentHistoryBestEffort] = struct{}{} + if db.mode == WritingFileWithHistoryEnforced { + if walHead.Features == nil { + walHead.Features = make(map[string]struct{}) } + walHead.Features[FeatureDeploymentHistoryBestEffort] = struct{}{} } if err := appendJSONLine(db.walFile, walHead); err != nil { // Remove the WAL created just above: O_EXCL means it is ours, and diff --git a/bundle/direct/dstate/state_test.go b/bundle/direct/dstate/state_test.go index 2395ee349a2..ce41c2dc7d1 100644 --- a/bundle/direct/dstate/state_test.go +++ b/bundle/direct/dstate/state_test.go @@ -2,7 +2,6 @@ package dstate import ( "encoding/json" - "fmt" "os" "path/filepath" "testing" @@ -422,24 +421,3 @@ func TestOpenFailureLeavesStateClosed(t *testing.T) { assert.Equal(t, "test-lineage", db.Data.Lineage) mustFinalize(t, &db) } - -func TestFileWriteErrorsRemainFatalWithHistory(t *testing.T) { - for _, mode := range []Mode{WritingFile, WritingFileWithHistory, WritingFileWithHistoryEnforced} { - t.Run(fmt.Sprint(mode), func(t *testing.T) { - path := filepath.Join(t.TempDir(), "resources.json") - data := NewDatabase("lineage", 1) - data.State["resources.jobs.job"] = ResourceEntry{ID: "old-id"} - var db DeploymentState - db.OpenWithData(path, data) - db.mode = mode - wal, err := os.Create(path + walSuffix) - require.NoError(t, err) - require.NoError(t, wal.Close()) - db.walFile = wal - require.Error(t, db.SaveState(t.Context(), "resources.jobs.job", "new-id", struct{}{}, nil)) - assert.Equal(t, "old-id", db.GetResourceID("resources.jobs.job")) - require.Error(t, db.DeleteState(t.Context(), "resources.jobs.job", false)) - assert.Equal(t, "old-id", db.GetResourceID("resources.jobs.job")) - }) - } -} From d61e89b936db63dcc7a224d3756f30bc90c4c425 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 11:09:25 +0000 Subject: [PATCH 36/52] Remove defensive panic when recording after drain --- libs/dms/operation_buffer.go | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index f1e419c58c0..9bea02cd3f7 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -96,16 +96,13 @@ func (b *OperationBuffer) RecordFailure(resourceKey, resourceID string, state js // record queues update for resourceKey. keepExistingUpdateIfAny (a failure) merges onto the // resource's last recorded state if one exists; otherwise (a state write) update is saved as -// that latest state. latestState and pending are updated under the one lock. Panics after Drain. +// that latest state. latestState and pending are updated under the one lock. func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, keepExistingUpdateIfAny bool, cause error) { b.mu.Lock() defer b.mu.Unlock() if b.bestEffort && b.err != nil { return } - if b.closed { - panic("recording operation after drain") - } if keepExistingUpdateIfAny { if latest, recorded := b.latestState[resourceKey]; recorded { update = latest.Merge(failureUpdate(cause)) From ad5ce69d54faf78b3595deb441e61741fce130f9 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 11:17:05 +0000 Subject: [PATCH 37/52] Simplify unbounded queue coordination with a condition variable --- libs/dms/operation_buffer.go | 69 ++++++++++++++----------------- libs/dms/operation_buffer_test.go | 19 ++++++++- 2 files changed, 48 insertions(+), 40 deletions(-) diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 9bea02cd3f7..b06b1a32647 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -22,7 +22,7 @@ type OperationBuffer struct { deploymentID string versionNum int - wake chan struct{} + ready *sync.Cond done chan struct{} cancel context.CancelFunc bestEffort bool @@ -40,8 +40,8 @@ type OperationBuffer struct { mu sync.Mutex // queue is unbounded and contains one key per pending update. - queue []string - closed bool + queue []string + draining bool // pending holds the newest update per resource key, absent once the writer takes it. pending map[string]OperationUpdate @@ -62,7 +62,6 @@ func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum i b := &OperationBuffer{ deploymentID: deploymentID, versionNum: versionNum, - wake: make(chan struct{}, 1), done: make(chan struct{}), cancel: cancel, bestEffort: options.BestEffort, @@ -70,6 +69,7 @@ func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum i latestState: make(map[string]OperationUpdate), sequenceIDs: make(map[string]int64), } + b.ready = sync.NewCond(&b.mu) go b.run(ctx) return b } @@ -122,43 +122,28 @@ func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, kee // Only when the resource is not already queued, enqueue it. if !queued { b.queue = append(b.queue, resourceKey) - b.notify() + b.ready.Signal() } } -// notify wakes the writer without blocking producers. -func (b *OperationBuffer) notify() { - select { - case b.wake <- struct{}{}: - default: - } -} - -// next waits for a queued key, drain, or cancellation. +// next waits until work is queued, draining begins, or the context is canceled. func (b *OperationBuffer) next(ctx context.Context) (string, bool) { - for { - b.mu.Lock() - if len(b.queue) > 0 { - key := b.queue[0] - b.queue[0] = "" - b.queue = b.queue[1:] - if len(b.queue) == 0 { - b.queue = nil - } - b.mu.Unlock() - return key, true - } - closed := b.closed - b.mu.Unlock() - if closed { - return "", false - } - select { - case <-b.wake: - case <-ctx.Done(): - return "", false - } + b.mu.Lock() + defer b.mu.Unlock() + + for len(b.queue) == 0 && !b.draining && ctx.Err() == nil { + b.ready.Wait() + } + if len(b.queue) == 0 { + return "", false } + key := b.queue[0] + b.queue[0] = "" + b.queue = b.queue[1:] + if len(b.queue) == 0 { + b.queue = nil + } + return key, true } // take claims the update waiting for resourceKey. @@ -175,6 +160,14 @@ func (b *OperationBuffer) run(ctx context.Context) { defer close(b.done) defer b.cancel() + // Wake an idle worker when its context is canceled. + stopWake := context.AfterFunc(ctx, func() { + b.mu.Lock() + b.ready.Signal() + b.mu.Unlock() + }) + defer stopWake() + for { resourceKey, ok := b.next(ctx) if err := ctx.Err(); err != nil { @@ -222,9 +215,9 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio // It returns the first recording error and is safe to call twice. func (b *OperationBuffer) Drain(ctx context.Context) error { b.mu.Lock() - b.closed = true + b.draining = true + b.ready.Signal() b.mu.Unlock() - b.notify() select { case <-b.done: return b.Err() diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index 117e2d4707e..c3343a466d7 100644 --- a/libs/dms/operation_buffer_test.go +++ b/libs/dms/operation_buffer_test.go @@ -1,10 +1,12 @@ package dms import ( + "context" "encoding/json" "errors" "fmt" "strings" + "sync" "testing" "testing/synctest" @@ -20,12 +22,13 @@ import ( // queued builds a buffer whose queue nothing drains, so a test drives record and take itself. func queued() *OperationBuffer { - return &OperationBuffer{ - wake: make(chan struct{}, 1), + b := &OperationBuffer{ cancel: func() {}, pending: make(map[string]OperationUpdate), latestState: make(map[string]OperationUpdate), } + b.ready = sync.NewCond(&b.mu) + return b } func stateUpdate(t *testing.T, name string) OperationUpdate { @@ -135,6 +138,18 @@ func TestBufferQueuesWithoutBackpressure(t *testing.T) { }) } +func TestBufferStopsIdleWriterOnCancellation(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + b := StartOperationBuffer(ctx, "dep-1", 1, BufferOptions{}) + synctest.Wait() + cancel() + <-b.done + assert.ErrorIs(t, b.Err(), context.Canceled) + }) +} + func TestBestEffortBufferStopsIdleWriterAfterError(t *testing.T) { synctest.Test(t, func(t *testing.T) { b := StartOperationBuffer(t.Context(), "dep-1", 1, BufferOptions{BestEffort: true}) From 22b54a5d9b3df9c0fd903f338f3b85f73330b9a5 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 11:32:23 +0000 Subject: [PATCH 38/52] Simplify operation buffering with 100ms idle polling --- .../bundle/dms/failed-recreate/output.txt | 42 +----------- acceptance/bundle/dms/failed-recreate/script | 4 +- .../operation-upload-fails-once/output.txt | 42 +----------- .../dms/operation-upload-fails-once/script | 8 +-- .../dms/operation-upload-fails-once/test.toml | 5 ++ .../bundle/dms/successful-recreate/output.txt | 64 ++----------------- .../bundle/dms/successful-recreate/script | 11 ++-- bundle/direct/dstate/mode_test.go | 35 ++++++++++ libs/dms/operation_buffer.go | 56 ++++++---------- libs/dms/operation_buffer_test.go | 15 +---- 10 files changed, 82 insertions(+), 200 deletions(-) diff --git a/acceptance/bundle/dms/failed-recreate/output.txt b/acceptance/bundle/dms/failed-recreate/output.txt index 860460c3375..eea5d4479ae 100644 --- a/acceptance/bundle/dms/failed-recreate/output.txt +++ b/acceptance/bundle/dms/failed-recreate/output.txt @@ -6,7 +6,7 @@ Created schemas.foo Files: 2 uploaded, 0 deleted Resources: 1 created, 0 changed, 0 deleted, 0 unchanged ->>> print_requests.py --dms //api/2.0/bundle +>>> print_requests.py --dms //api/2.0/bundle --method POST { "method": "POST", "path": "/api/2.0/bundle/deployments", @@ -37,20 +37,6 @@ Resources: 1 created, 0 changed, 0 deleted, 0 unchanged "version_type": "VERSION_TYPE_DEPLOY" } } -{ - "method": "PATCH", - "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/1/operations/schemas.foo", - "q": { - "update_mask": "state,error_message,resource_id,status" - }, - "body": { - "error_message": "", - "resource_id": "main.dms_failed_recreate_schema_[UNIQUE_NAME]", - "sequence_id": 0, - "state": "{\"state\":{\"catalog_name\":\"main\",\"comment\":\"v1\",\"name\":\"dms_failed_recreate_schema_[UNIQUE_NAME]\"}}", - "status": "OPERATION_STATUS_SUCCEEDED" - } -} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/1/complete", @@ -116,7 +102,7 @@ API message: Fault injected by test. Files: 2 uploaded, 0 deleted ->>> print_requests.py --dms //api/2.0/bundle +>>> print_requests.py --dms //api/2.0/bundle --method POST { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions", @@ -135,30 +121,6 @@ Files: 2 uploaded, 0 deleted "version_type": "VERSION_TYPE_DEPLOY" } } -{ - "method": "PATCH", - "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", - "q": { - "update_mask": "state,error_message,status" - }, - "body": { - "error_message": "", - "sequence_id": 0, - "status": "OPERATION_STATUS_PENDING" - } -} -{ - "method": "PATCH", - "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", - "q": { - "update_mask": "state,error_message,status" - }, - "body": { - "error_message": "Fault injected by test. (400 INVALID_PARAMETER_VALUE)", - "sequence_id": 1, - "status": "OPERATION_STATUS_FAILED" - } -} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/complete", diff --git a/acceptance/bundle/dms/failed-recreate/script b/acceptance/bundle/dms/failed-recreate/script index 60bcc99df10..3cbcf917d30 100644 --- a/acceptance/bundle/dms/failed-recreate/script +++ b/acceptance/bundle/dms/failed-recreate/script @@ -2,7 +2,7 @@ envsubst < databricks.yml.tmpl > databricks.yml title "Deploy the schema, so there is a resource to recreate" trace $CLI bundle deploy -trace print_requests.py --dms //api/2.0/bundle +trace print_requests.py --dms //api/2.0/bundle --method POST title "The plan is a no-op and records the deployment_history feature" dms_plan @@ -25,7 +25,7 @@ trace update_file.py databricks.yml "catalog_name: main" "catalog_name: ${other_ trace fault.py "POST /api/2.1/unity-catalog/schemas" 400 0 1 INVALID_PARAMETER_VALUE trace musterr $CLI bundle deploy --auto-approve # Not sorted: the point of this test is the order the operations are recorded in. -trace print_requests.py --dms //api/2.0/bundle +trace print_requests.py --dms //api/2.0/bundle --method POST title "The resource is not listed: state is what projects a resource, and the failed recreate left none" deployment_id=$(MSYS_NO_PATHCONV=1 $CLI workspace get-status "/Workspace/Users/${CURRENT_USER_NAME}/.bundle/dms-failed-recreate-${UNIQUE_NAME}/default/state/resources.deployment.json" | python3 -c 'import sys, json; print(json.load(sys.stdin)["object_id"])') diff --git a/acceptance/bundle/dms/operation-upload-fails-once/output.txt b/acceptance/bundle/dms/operation-upload-fails-once/output.txt index dc976dadead..67c97aef3fc 100644 --- a/acceptance/bundle/dms/operation-upload-fails-once/output.txt +++ b/acceptance/bundle/dms/operation-upload-fails-once/output.txt @@ -44,7 +44,7 @@ Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged "state": {} } -=== A recreate records twice, and the first one fails: the second must still carry the sequence id staging left, because a rejected update earned no new one +=== The first uploaded update fails; the version still completes with failure >>> fault.py PATCH /api/2.0/bundle/* 500 0 1 >>> update_file.py databricks.yml catalog_name: main catalog_name: dms_other_[UNIQUE_NAME] @@ -57,23 +57,13 @@ This action will result in the deletion or recreation of the following UC schema Recreated schemas.foo Error: recording operation for resources.schemas.foo: Fault injected by test. (500 INJECTED) -Endpoint: PATCH [DATABRICKS_URL]/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo?update_mask=state%2Cerror_message%2Cstatus +Endpoint: PATCH [DATABRICKS_URL]/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo?update_mask=[UPDATE_MASK] HTTP Status: 500 Internal Server Error API error_code: INJECTED API message: Fault injected by test. Files: 2 uploaded, 0 deleted -=== The service lists the schema with the failure ->>> MSYS_NO_PATHCONV=1 [CLI] bundle-deployments list-resources deployments/[DEPLOYMENT_ID] -o json -[ - { - "last_action_type": "OPERATION_ACTION_TYPE_RECREATE", - "last_version_id": "2", - "resource_key": "schemas.foo" - } -] - === The raw recorded state after the failure >>> print_state.py --no-dms { @@ -87,7 +77,7 @@ Files: 2 uploaded, 0 deleted "state": {} } ->>> print_requests.py --dms //api/2.0/bundle +>>> print_requests.py --dms //api/2.0/bundle --method POST { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions", @@ -106,32 +96,6 @@ Files: 2 uploaded, 0 deleted "version_type": "VERSION_TYPE_DEPLOY" } } -{ - "method": "PATCH", - "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", - "q": { - "update_mask": "state,error_message,status" - }, - "body": { - "error_message": "", - "sequence_id": 0, - "status": "OPERATION_STATUS_PENDING" - } -} -{ - "method": "PATCH", - "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", - "q": { - "update_mask": "state,error_message,resource_id,status" - }, - "body": { - "error_message": "", - "resource_id": "dms_other_[UNIQUE_NAME].dms_upload_fails_once_schema_[UNIQUE_NAME]", - "sequence_id": 0, - "state": "{\"state\":{\"catalog_name\":\"dms_other_[UNIQUE_NAME]\",\"comment\":\"v1\",\"name\":\"dms_upload_fails_once_schema_[UNIQUE_NAME]\"}}", - "status": "OPERATION_STATUS_SUCCEEDED" - } -} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/complete", diff --git a/acceptance/bundle/dms/operation-upload-fails-once/script b/acceptance/bundle/dms/operation-upload-fails-once/script index 2b539134571..885fec058e4 100644 --- a/acceptance/bundle/dms/operation-upload-fails-once/script +++ b/acceptance/bundle/dms/operation-upload-fails-once/script @@ -15,7 +15,7 @@ trace MSYS_NO_PATHCONV=1 $CLI bundle-deployments list-resources "deployments/${d title "The raw recorded state after deploy" trace print_state.py --no-dms -title "A recreate records twice, and the first one fails: the second must still carry the sequence id staging left, because a rejected update earned no new one" +title "The first uploaded update fails; the version still completes with failure" trace fault.py "PATCH /api/2.0/bundle/*" 500 0 1 # Unity Catalog has no catalog called "other", so make one to move the schema into. other_catalog="dms_other_${UNIQUE_NAME}" @@ -24,11 +24,7 @@ trap '$CLI catalogs delete "$other_catalog" --force >> LOG.catalog 2>&1' EXIT trace update_file.py databricks.yml "catalog_name: main" "catalog_name: ${other_catalog}" trace musterr $CLI bundle deploy --auto-approve -title "The service lists the schema with the failure" -deployment_id=$(MSYS_NO_PATHCONV=1 $CLI workspace get-status "/Workspace/Users/${CURRENT_USER_NAME}/.bundle/dms-upload-fails-once-${UNIQUE_NAME}/default/state/resources.deployment.json" | python3 -c 'import sys, json; print(json.load(sys.stdin)["object_id"])') -trace MSYS_NO_PATHCONV=1 $CLI bundle-deployments list-resources "deployments/${deployment_id}" -o json | jq -S 'map({resource_key, last_version_id, last_action_type})' - title "The raw recorded state after the failure" trace print_state.py --no-dms -trace print_requests.py --dms //api/2.0/bundle +trace print_requests.py --dms //api/2.0/bundle --method POST diff --git a/acceptance/bundle/dms/operation-upload-fails-once/test.toml b/acceptance/bundle/dms/operation-upload-fails-once/test.toml index 5baa9c9d923..49da31fbf5b 100644 --- a/acceptance/bundle/dms/operation-upload-fails-once/test.toml +++ b/acceptance/bundle/dms/operation-upload-fails-once/test.toml @@ -1,3 +1,8 @@ # Local only: this case needs an injected failure - DMS itself must reject the operation upload, once - and fault.py registers # its rule on the fake, which a cloud run does not have. Cloud = false + +# Coalescing can make either the pending or final state the first rejected update. +[[Repls]] +Old = '\?update_mask=\S+' +New = '?update_mask=[UPDATE_MASK]' diff --git a/acceptance/bundle/dms/successful-recreate/output.txt b/acceptance/bundle/dms/successful-recreate/output.txt index f2ecc3f8cbc..fa7cba315cd 100644 --- a/acceptance/bundle/dms/successful-recreate/output.txt +++ b/acceptance/bundle/dms/successful-recreate/output.txt @@ -6,7 +6,7 @@ Created schemas.foo Files: 2 uploaded, 0 deleted Resources: 1 created, 0 changed, 0 deleted, 0 unchanged ->>> print_requests.py --dms //api/2.0/bundle +>>> print_requests.py --dms //api/2.0/bundle --method POST { "method": "POST", "path": "/api/2.0/bundle/deployments", @@ -37,20 +37,6 @@ Resources: 1 created, 0 changed, 0 deleted, 0 unchanged "version_type": "VERSION_TYPE_DEPLOY" } } -{ - "method": "PATCH", - "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/1/operations/schemas.foo", - "q": { - "update_mask": "state,error_message,resource_id,status" - }, - "body": { - "error_message": "", - "resource_id": "main.dms_successful_recreate_schema_[UNIQUE_NAME]", - "sequence_id": 0, - "state": "{\"state\":{\"catalog_name\":\"main\",\"comment\":\"v1\",\"name\":\"dms_successful_recreate_schema_[UNIQUE_NAME]\"}}", - "status": "OPERATION_STATUS_SUCCEEDED" - } -} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/1/complete", @@ -97,7 +83,7 @@ Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged "state": {} } -=== Recreate writes state twice - the entry is dropped, then the new resource is saved +=== Recreate queues the delete and create state; pending updates may coalesce >>> update_file.py databricks.yml catalog_name: main catalog_name: dms_other_[UNIQUE_NAME] >>> [CLI] bundle deploy --auto-approve @@ -109,7 +95,7 @@ Recreated schemas.foo Files: 2 uploaded, 0 deleted Resources: 1 created, 0 changed, 1 deleted, 0 unchanged ->>> print_requests.py --dms //api/2.0/bundle +>>> print_requests.py --dms //api/2.0/bundle --method POST { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions", @@ -128,32 +114,6 @@ Resources: 1 created, 0 changed, 1 deleted, 0 unchanged "version_type": "VERSION_TYPE_DEPLOY" } } -{ - "method": "PATCH", - "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", - "q": { - "update_mask": "state,error_message,status" - }, - "body": { - "error_message": "", - "sequence_id": 0, - "status": "OPERATION_STATUS_PENDING" - } -} -{ - "method": "PATCH", - "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", - "q": { - "update_mask": "state,error_message,resource_id,status" - }, - "body": { - "error_message": "", - "resource_id": "dms_other_[UNIQUE_NAME].dms_successful_recreate_schema_[UNIQUE_NAME]", - "sequence_id": 1, - "state": "{\"state\":{\"catalog_name\":\"dms_other_[UNIQUE_NAME]\",\"comment\":\"v1\",\"name\":\"dms_successful_recreate_schema_[UNIQUE_NAME]\"}}", - "status": "OPERATION_STATUS_SUCCEEDED" - } -} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/complete", @@ -212,7 +172,7 @@ All files and directories at the following location will be deleted: /Workspace/ Destroy: 1 deleted ->>> print_requests.py --dms //api/2.0/bundle +>>> print_requests.py --dms //api/2.0/bundle --method POST { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions", @@ -231,18 +191,6 @@ Destroy: 1 deleted "version_type": "VERSION_TYPE_DESTROY" } } -{ - "method": "PATCH", - "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/3/operations/schemas.foo", - "q": { - "update_mask": "state,error_message,status" - }, - "body": { - "error_message": "", - "sequence_id": 0, - "status": "OPERATION_STATUS_SUCCEEDED" - } -} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/3/complete", @@ -250,7 +198,3 @@ Destroy: 1 deleted "completion_reason": "VERSION_COMPLETE_SUCCESS" } } -{ - "method": "DELETE", - "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]" -} diff --git a/acceptance/bundle/dms/successful-recreate/script b/acceptance/bundle/dms/successful-recreate/script index 87bdc98bc34..eb20654a151 100644 --- a/acceptance/bundle/dms/successful-recreate/script +++ b/acceptance/bundle/dms/successful-recreate/script @@ -2,7 +2,7 @@ envsubst < databricks.yml.tmpl > databricks.yml title "Deploy: the state write records the resource" trace $CLI bundle deploy -trace print_requests.py --dms //api/2.0/bundle +trace print_requests.py --dms //api/2.0/bundle --method POST title "The plan is a no-op now and records the deployment_history feature" dms_plan @@ -15,16 +15,15 @@ trace MSYS_NO_PATHCONV=1 $CLI bundle-deployments list-resources "deployments/${d title "The raw recorded state the service holds for this deployment" trace print_state.py --no-dms -title "Recreate writes state twice - the entry is dropped, then the new resource is saved" -# One operation per resource per version; both writes land on the same one. -# The drop opens it IN_PROGRESS, the save patches it to SUCCEEDED. +title "Recreate queues the delete and create state; pending updates may coalesce" +# Both state writes target the same operation and may be combined before upload. # Unity Catalog has no catalog called "other", so make one to move the schema into. other_catalog="dms_other_${UNIQUE_NAME}" $CLI catalogs create "$other_catalog" &> LOG.catalog trap '$CLI catalogs delete "$other_catalog" --force >> LOG.catalog 2>&1' EXIT trace update_file.py databricks.yml "catalog_name: main" "catalog_name: ${other_catalog}" trace $CLI bundle deploy --auto-approve -trace print_requests.py --dms //api/2.0/bundle +trace print_requests.py --dms //api/2.0/bundle --method POST title "The recreate is done: nothing left to plan, and the serial moved to 2" dms_plan @@ -37,5 +36,5 @@ trace print_state.py --no-dms title "Destroy: the delete is recorded with the id and no state" trace $CLI bundle destroy --auto-approve -trace print_requests.py --dms //api/2.0/bundle +trace print_requests.py --dms //api/2.0/bundle --method POST rm -f "$OUT_REQUESTS" diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index 670adeb043a..edc2dafbc1b 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -263,6 +263,41 @@ func TestUpgradeToWrite(t *testing.T) { } } +func TestDMSRecordingContinuesAfterUploadError(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + firstRequest := make(chan struct{}) + updates := 0 + ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + var body struct { + SequenceID int64 `json:"sequence_id"` + State string `json:"state"` + } + assert.NoError(t, json.NewDecoder(r.Body).Decode(&body)) + assert.Zero(t, body.SequenceID, "a rejected update earns no sequence ID") + updates++ + if updates == 1 { + close(firstRequest) + http.Error(w, `{"error_code":"INVALID_PARAMETER_VALUE","message":"first update failed"}`, http.StatusBadRequest) + return + } + assert.JSONEq(t, `{"state":{"name":"second"}}`, body.State) + _, _ = w.Write([]byte(`{"sequence_id":1}`)) + }) + var db dstate.DeploymentState + path := filepath.Join(t.TempDir(), "resources.json") + require.NoError(t, db.Open(ctx, path, dstate.WritingDMS, false, dstate.OpenDmsArgs{})) + db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) + require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", map[string]string{"name": "first"}, nil)) + <-firstRequest + synctest.Wait() + require.ErrorContains(t, db.RecordingError(), "first update failed") + require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", map[string]string{"name": "second"}, nil)) + _, err := db.Finalize(ctx) + require.ErrorContains(t, err, "first update failed") + assert.Equal(t, 2, updates) + }) +} + func TestHistoryErrorsFailOpen(t *testing.T) { for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { t.Run(fmt.Sprint(mode), func(t *testing.T) { diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index b06b1a32647..2128be2e5ca 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "sync" + "time" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/databricks-sdk-go/service/bundledeployments" @@ -22,7 +23,6 @@ type OperationBuffer struct { deploymentID string versionNum int - ready *sync.Cond done chan struct{} cancel context.CancelFunc bestEffort bool @@ -69,7 +69,6 @@ func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum i latestState: make(map[string]OperationUpdate), sequenceIDs: make(map[string]int64), } - b.ready = sync.NewCond(&b.mu) go b.run(ctx) return b } @@ -122,30 +121,9 @@ func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, kee // Only when the resource is not already queued, enqueue it. if !queued { b.queue = append(b.queue, resourceKey) - b.ready.Signal() } } -// next waits until work is queued, draining begins, or the context is canceled. -func (b *OperationBuffer) next(ctx context.Context) (string, bool) { - b.mu.Lock() - defer b.mu.Unlock() - - for len(b.queue) == 0 && !b.draining && ctx.Err() == nil { - b.ready.Wait() - } - if len(b.queue) == 0 { - return "", false - } - key := b.queue[0] - b.queue[0] = "" - b.queue = b.queue[1:] - if len(b.queue) == 0 { - b.queue = nil - } - return key, true -} - // take claims the update waiting for resourceKey. func (b *OperationBuffer) take(resourceKey string) (OperationUpdate, bool) { b.mu.Lock() @@ -160,23 +138,32 @@ func (b *OperationBuffer) run(ctx context.Context) { defer close(b.done) defer b.cancel() - // Wake an idle worker when its context is canceled. - stopWake := context.AfterFunc(ctx, func() { - b.mu.Lock() - b.ready.Signal() - b.mu.Unlock() - }) - defer stopWake() - for { - resourceKey, ok := b.next(ctx) if err := ctx.Err(); err != nil { b.setErr(err) return } - if !ok { - return + b.mu.Lock() + if len(b.queue) == 0 { + draining := b.draining + b.mu.Unlock() + if draining { + return + } + // Poll an empty queue without holding the lock or delaying cancellation. + select { + case <-time.After(100 * time.Millisecond): + case <-ctx.Done(): + } + continue } + resourceKey := b.queue[0] + b.queue[0] = "" + b.queue = b.queue[1:] + if len(b.queue) == 0 { + b.queue = nil + } + b.mu.Unlock() update, ok := b.take(resourceKey) if !ok { // Unreachable: a key is queued only when nothing was waiting for it. Guard so @@ -216,7 +203,6 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio func (b *OperationBuffer) Drain(ctx context.Context) error { b.mu.Lock() b.draining = true - b.ready.Signal() b.mu.Unlock() select { case <-b.done: diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index c3343a466d7..5c8d1c53bf6 100644 --- a/libs/dms/operation_buffer_test.go +++ b/libs/dms/operation_buffer_test.go @@ -6,7 +6,6 @@ import ( "errors" "fmt" "strings" - "sync" "testing" "testing/synctest" @@ -22,13 +21,11 @@ import ( // queued builds a buffer whose queue nothing drains, so a test drives record and take itself. func queued() *OperationBuffer { - b := &OperationBuffer{ + return &OperationBuffer{ cancel: func() {}, pending: make(map[string]OperationUpdate), latestState: make(map[string]OperationUpdate), } - b.ready = sync.NewCond(&b.mu) - return b } func stateUpdate(t *testing.T, name string) OperationUpdate { @@ -74,7 +71,7 @@ func TestBufferFailsOnOversizedState(t *testing.T) { } func TestDrainIsIdempotent(t *testing.T) { - // An idle writer must wake up when the queue is closed. + // An idle writer must finish when draining begins. b := StartOperationBuffer(t.Context(), "dep-1", 1, BufferOptions{}) require.NoError(t, b.Drain(t.Context())) require.NoError(t, b.Drain(t.Context())) @@ -126,15 +123,9 @@ func TestBufferQueuesWithoutBackpressure(t *testing.T) { assert.NoError(t, b.Err()) assert.Len(t, b.queue, operationCount) assert.Len(t, b.pending, operationCount) - for i := range operationCount { - key, ok := b.next(t.Context()) - require.True(t, ok) + for i, key := range b.queue { assert.Equal(t, fmt.Sprintf("resources.jobs.job_%d", i), key) - _, ok = b.take(key) - require.True(t, ok) } - assert.Empty(t, b.queue) - assert.Empty(t, b.pending) }) } From 2333f443cb960f02375cda94a242b612409e8a0c Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 11:41:48 +0000 Subject: [PATCH 39/52] Preserve existing DMS behavior and restore acceptance coverage --- .../bundle/dms/failed-recreate/output.txt | 42 ++++++++- acceptance/bundle/dms/failed-recreate/script | 4 +- .../operation-upload-fails-once/output.txt | 42 ++++++++- .../dms/operation-upload-fails-once/script | 8 +- .../dms/operation-upload-fails-once/test.toml | 5 -- .../bundle/dms/successful-recreate/output.txt | 64 ++++++++++++- .../bundle/dms/successful-recreate/script | 11 +-- bundle/direct/dstate/mode_test.go | 35 -------- libs/dms/operation_buffer.go | 90 +++++++++++++------ libs/dms/operation_buffer_test.go | 10 ++- 10 files changed, 221 insertions(+), 90 deletions(-) diff --git a/acceptance/bundle/dms/failed-recreate/output.txt b/acceptance/bundle/dms/failed-recreate/output.txt index eea5d4479ae..860460c3375 100644 --- a/acceptance/bundle/dms/failed-recreate/output.txt +++ b/acceptance/bundle/dms/failed-recreate/output.txt @@ -6,7 +6,7 @@ Created schemas.foo Files: 2 uploaded, 0 deleted Resources: 1 created, 0 changed, 0 deleted, 0 unchanged ->>> print_requests.py --dms //api/2.0/bundle --method POST +>>> print_requests.py --dms //api/2.0/bundle { "method": "POST", "path": "/api/2.0/bundle/deployments", @@ -37,6 +37,20 @@ Resources: 1 created, 0 changed, 0 deleted, 0 unchanged "version_type": "VERSION_TYPE_DEPLOY" } } +{ + "method": "PATCH", + "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/1/operations/schemas.foo", + "q": { + "update_mask": "state,error_message,resource_id,status" + }, + "body": { + "error_message": "", + "resource_id": "main.dms_failed_recreate_schema_[UNIQUE_NAME]", + "sequence_id": 0, + "state": "{\"state\":{\"catalog_name\":\"main\",\"comment\":\"v1\",\"name\":\"dms_failed_recreate_schema_[UNIQUE_NAME]\"}}", + "status": "OPERATION_STATUS_SUCCEEDED" + } +} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/1/complete", @@ -102,7 +116,7 @@ API message: Fault injected by test. Files: 2 uploaded, 0 deleted ->>> print_requests.py --dms //api/2.0/bundle --method POST +>>> print_requests.py --dms //api/2.0/bundle { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions", @@ -121,6 +135,30 @@ Files: 2 uploaded, 0 deleted "version_type": "VERSION_TYPE_DEPLOY" } } +{ + "method": "PATCH", + "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", + "q": { + "update_mask": "state,error_message,status" + }, + "body": { + "error_message": "", + "sequence_id": 0, + "status": "OPERATION_STATUS_PENDING" + } +} +{ + "method": "PATCH", + "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", + "q": { + "update_mask": "state,error_message,status" + }, + "body": { + "error_message": "Fault injected by test. (400 INVALID_PARAMETER_VALUE)", + "sequence_id": 1, + "status": "OPERATION_STATUS_FAILED" + } +} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/complete", diff --git a/acceptance/bundle/dms/failed-recreate/script b/acceptance/bundle/dms/failed-recreate/script index 3cbcf917d30..60bcc99df10 100644 --- a/acceptance/bundle/dms/failed-recreate/script +++ b/acceptance/bundle/dms/failed-recreate/script @@ -2,7 +2,7 @@ envsubst < databricks.yml.tmpl > databricks.yml title "Deploy the schema, so there is a resource to recreate" trace $CLI bundle deploy -trace print_requests.py --dms //api/2.0/bundle --method POST +trace print_requests.py --dms //api/2.0/bundle title "The plan is a no-op and records the deployment_history feature" dms_plan @@ -25,7 +25,7 @@ trace update_file.py databricks.yml "catalog_name: main" "catalog_name: ${other_ trace fault.py "POST /api/2.1/unity-catalog/schemas" 400 0 1 INVALID_PARAMETER_VALUE trace musterr $CLI bundle deploy --auto-approve # Not sorted: the point of this test is the order the operations are recorded in. -trace print_requests.py --dms //api/2.0/bundle --method POST +trace print_requests.py --dms //api/2.0/bundle title "The resource is not listed: state is what projects a resource, and the failed recreate left none" deployment_id=$(MSYS_NO_PATHCONV=1 $CLI workspace get-status "/Workspace/Users/${CURRENT_USER_NAME}/.bundle/dms-failed-recreate-${UNIQUE_NAME}/default/state/resources.deployment.json" | python3 -c 'import sys, json; print(json.load(sys.stdin)["object_id"])') diff --git a/acceptance/bundle/dms/operation-upload-fails-once/output.txt b/acceptance/bundle/dms/operation-upload-fails-once/output.txt index 67c97aef3fc..dc976dadead 100644 --- a/acceptance/bundle/dms/operation-upload-fails-once/output.txt +++ b/acceptance/bundle/dms/operation-upload-fails-once/output.txt @@ -44,7 +44,7 @@ Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged "state": {} } -=== The first uploaded update fails; the version still completes with failure +=== A recreate records twice, and the first one fails: the second must still carry the sequence id staging left, because a rejected update earned no new one >>> fault.py PATCH /api/2.0/bundle/* 500 0 1 >>> update_file.py databricks.yml catalog_name: main catalog_name: dms_other_[UNIQUE_NAME] @@ -57,13 +57,23 @@ This action will result in the deletion or recreation of the following UC schema Recreated schemas.foo Error: recording operation for resources.schemas.foo: Fault injected by test. (500 INJECTED) -Endpoint: PATCH [DATABRICKS_URL]/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo?update_mask=[UPDATE_MASK] +Endpoint: PATCH [DATABRICKS_URL]/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo?update_mask=state%2Cerror_message%2Cstatus HTTP Status: 500 Internal Server Error API error_code: INJECTED API message: Fault injected by test. Files: 2 uploaded, 0 deleted +=== The service lists the schema with the failure +>>> MSYS_NO_PATHCONV=1 [CLI] bundle-deployments list-resources deployments/[DEPLOYMENT_ID] -o json +[ + { + "last_action_type": "OPERATION_ACTION_TYPE_RECREATE", + "last_version_id": "2", + "resource_key": "schemas.foo" + } +] + === The raw recorded state after the failure >>> print_state.py --no-dms { @@ -77,7 +87,7 @@ Files: 2 uploaded, 0 deleted "state": {} } ->>> print_requests.py --dms //api/2.0/bundle --method POST +>>> print_requests.py --dms //api/2.0/bundle { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions", @@ -96,6 +106,32 @@ Files: 2 uploaded, 0 deleted "version_type": "VERSION_TYPE_DEPLOY" } } +{ + "method": "PATCH", + "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", + "q": { + "update_mask": "state,error_message,status" + }, + "body": { + "error_message": "", + "sequence_id": 0, + "status": "OPERATION_STATUS_PENDING" + } +} +{ + "method": "PATCH", + "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", + "q": { + "update_mask": "state,error_message,resource_id,status" + }, + "body": { + "error_message": "", + "resource_id": "dms_other_[UNIQUE_NAME].dms_upload_fails_once_schema_[UNIQUE_NAME]", + "sequence_id": 0, + "state": "{\"state\":{\"catalog_name\":\"dms_other_[UNIQUE_NAME]\",\"comment\":\"v1\",\"name\":\"dms_upload_fails_once_schema_[UNIQUE_NAME]\"}}", + "status": "OPERATION_STATUS_SUCCEEDED" + } +} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/complete", diff --git a/acceptance/bundle/dms/operation-upload-fails-once/script b/acceptance/bundle/dms/operation-upload-fails-once/script index 885fec058e4..2b539134571 100644 --- a/acceptance/bundle/dms/operation-upload-fails-once/script +++ b/acceptance/bundle/dms/operation-upload-fails-once/script @@ -15,7 +15,7 @@ trace MSYS_NO_PATHCONV=1 $CLI bundle-deployments list-resources "deployments/${d title "The raw recorded state after deploy" trace print_state.py --no-dms -title "The first uploaded update fails; the version still completes with failure" +title "A recreate records twice, and the first one fails: the second must still carry the sequence id staging left, because a rejected update earned no new one" trace fault.py "PATCH /api/2.0/bundle/*" 500 0 1 # Unity Catalog has no catalog called "other", so make one to move the schema into. other_catalog="dms_other_${UNIQUE_NAME}" @@ -24,7 +24,11 @@ trap '$CLI catalogs delete "$other_catalog" --force >> LOG.catalog 2>&1' EXIT trace update_file.py databricks.yml "catalog_name: main" "catalog_name: ${other_catalog}" trace musterr $CLI bundle deploy --auto-approve +title "The service lists the schema with the failure" +deployment_id=$(MSYS_NO_PATHCONV=1 $CLI workspace get-status "/Workspace/Users/${CURRENT_USER_NAME}/.bundle/dms-upload-fails-once-${UNIQUE_NAME}/default/state/resources.deployment.json" | python3 -c 'import sys, json; print(json.load(sys.stdin)["object_id"])') +trace MSYS_NO_PATHCONV=1 $CLI bundle-deployments list-resources "deployments/${deployment_id}" -o json | jq -S 'map({resource_key, last_version_id, last_action_type})' + title "The raw recorded state after the failure" trace print_state.py --no-dms -trace print_requests.py --dms //api/2.0/bundle --method POST +trace print_requests.py --dms //api/2.0/bundle diff --git a/acceptance/bundle/dms/operation-upload-fails-once/test.toml b/acceptance/bundle/dms/operation-upload-fails-once/test.toml index 49da31fbf5b..5baa9c9d923 100644 --- a/acceptance/bundle/dms/operation-upload-fails-once/test.toml +++ b/acceptance/bundle/dms/operation-upload-fails-once/test.toml @@ -1,8 +1,3 @@ # Local only: this case needs an injected failure - DMS itself must reject the operation upload, once - and fault.py registers # its rule on the fake, which a cloud run does not have. Cloud = false - -# Coalescing can make either the pending or final state the first rejected update. -[[Repls]] -Old = '\?update_mask=\S+' -New = '?update_mask=[UPDATE_MASK]' diff --git a/acceptance/bundle/dms/successful-recreate/output.txt b/acceptance/bundle/dms/successful-recreate/output.txt index fa7cba315cd..f2ecc3f8cbc 100644 --- a/acceptance/bundle/dms/successful-recreate/output.txt +++ b/acceptance/bundle/dms/successful-recreate/output.txt @@ -6,7 +6,7 @@ Created schemas.foo Files: 2 uploaded, 0 deleted Resources: 1 created, 0 changed, 0 deleted, 0 unchanged ->>> print_requests.py --dms //api/2.0/bundle --method POST +>>> print_requests.py --dms //api/2.0/bundle { "method": "POST", "path": "/api/2.0/bundle/deployments", @@ -37,6 +37,20 @@ Resources: 1 created, 0 changed, 0 deleted, 0 unchanged "version_type": "VERSION_TYPE_DEPLOY" } } +{ + "method": "PATCH", + "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/1/operations/schemas.foo", + "q": { + "update_mask": "state,error_message,resource_id,status" + }, + "body": { + "error_message": "", + "resource_id": "main.dms_successful_recreate_schema_[UNIQUE_NAME]", + "sequence_id": 0, + "state": "{\"state\":{\"catalog_name\":\"main\",\"comment\":\"v1\",\"name\":\"dms_successful_recreate_schema_[UNIQUE_NAME]\"}}", + "status": "OPERATION_STATUS_SUCCEEDED" + } +} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/1/complete", @@ -83,7 +97,7 @@ Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged "state": {} } -=== Recreate queues the delete and create state; pending updates may coalesce +=== Recreate writes state twice - the entry is dropped, then the new resource is saved >>> update_file.py databricks.yml catalog_name: main catalog_name: dms_other_[UNIQUE_NAME] >>> [CLI] bundle deploy --auto-approve @@ -95,7 +109,7 @@ Recreated schemas.foo Files: 2 uploaded, 0 deleted Resources: 1 created, 0 changed, 1 deleted, 0 unchanged ->>> print_requests.py --dms //api/2.0/bundle --method POST +>>> print_requests.py --dms //api/2.0/bundle { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions", @@ -114,6 +128,32 @@ Resources: 1 created, 0 changed, 1 deleted, 0 unchanged "version_type": "VERSION_TYPE_DEPLOY" } } +{ + "method": "PATCH", + "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", + "q": { + "update_mask": "state,error_message,status" + }, + "body": { + "error_message": "", + "sequence_id": 0, + "status": "OPERATION_STATUS_PENDING" + } +} +{ + "method": "PATCH", + "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/operations/schemas.foo", + "q": { + "update_mask": "state,error_message,resource_id,status" + }, + "body": { + "error_message": "", + "resource_id": "dms_other_[UNIQUE_NAME].dms_successful_recreate_schema_[UNIQUE_NAME]", + "sequence_id": 1, + "state": "{\"state\":{\"catalog_name\":\"dms_other_[UNIQUE_NAME]\",\"comment\":\"v1\",\"name\":\"dms_successful_recreate_schema_[UNIQUE_NAME]\"}}", + "status": "OPERATION_STATUS_SUCCEEDED" + } +} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/2/complete", @@ -172,7 +212,7 @@ All files and directories at the following location will be deleted: /Workspace/ Destroy: 1 deleted ->>> print_requests.py --dms //api/2.0/bundle --method POST +>>> print_requests.py --dms //api/2.0/bundle { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions", @@ -191,6 +231,18 @@ Destroy: 1 deleted "version_type": "VERSION_TYPE_DESTROY" } } +{ + "method": "PATCH", + "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/3/operations/schemas.foo", + "q": { + "update_mask": "state,error_message,status" + }, + "body": { + "error_message": "", + "sequence_id": 0, + "status": "OPERATION_STATUS_SUCCEEDED" + } +} { "method": "POST", "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]/versions/3/complete", @@ -198,3 +250,7 @@ Destroy: 1 deleted "completion_reason": "VERSION_COMPLETE_SUCCESS" } } +{ + "method": "DELETE", + "path": "/api/2.0/bundle/deployments/[DEPLOYMENT_ID]" +} diff --git a/acceptance/bundle/dms/successful-recreate/script b/acceptance/bundle/dms/successful-recreate/script index eb20654a151..87bdc98bc34 100644 --- a/acceptance/bundle/dms/successful-recreate/script +++ b/acceptance/bundle/dms/successful-recreate/script @@ -2,7 +2,7 @@ envsubst < databricks.yml.tmpl > databricks.yml title "Deploy: the state write records the resource" trace $CLI bundle deploy -trace print_requests.py --dms //api/2.0/bundle --method POST +trace print_requests.py --dms //api/2.0/bundle title "The plan is a no-op now and records the deployment_history feature" dms_plan @@ -15,15 +15,16 @@ trace MSYS_NO_PATHCONV=1 $CLI bundle-deployments list-resources "deployments/${d title "The raw recorded state the service holds for this deployment" trace print_state.py --no-dms -title "Recreate queues the delete and create state; pending updates may coalesce" -# Both state writes target the same operation and may be combined before upload. +title "Recreate writes state twice - the entry is dropped, then the new resource is saved" +# One operation per resource per version; both writes land on the same one. +# The drop opens it IN_PROGRESS, the save patches it to SUCCEEDED. # Unity Catalog has no catalog called "other", so make one to move the schema into. other_catalog="dms_other_${UNIQUE_NAME}" $CLI catalogs create "$other_catalog" &> LOG.catalog trap '$CLI catalogs delete "$other_catalog" --force >> LOG.catalog 2>&1' EXIT trace update_file.py databricks.yml "catalog_name: main" "catalog_name: ${other_catalog}" trace $CLI bundle deploy --auto-approve -trace print_requests.py --dms //api/2.0/bundle --method POST +trace print_requests.py --dms //api/2.0/bundle title "The recreate is done: nothing left to plan, and the serial moved to 2" dms_plan @@ -36,5 +37,5 @@ trace print_state.py --no-dms title "Destroy: the delete is recorded with the id and no state" trace $CLI bundle destroy --auto-approve -trace print_requests.py --dms //api/2.0/bundle --method POST +trace print_requests.py --dms //api/2.0/bundle rm -f "$OUT_REQUESTS" diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index edc2dafbc1b..670adeb043a 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -263,41 +263,6 @@ func TestUpgradeToWrite(t *testing.T) { } } -func TestDMSRecordingContinuesAfterUploadError(t *testing.T) { - synctest.Test(t, func(t *testing.T) { - firstRequest := make(chan struct{}) - updates := 0 - ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { - var body struct { - SequenceID int64 `json:"sequence_id"` - State string `json:"state"` - } - assert.NoError(t, json.NewDecoder(r.Body).Decode(&body)) - assert.Zero(t, body.SequenceID, "a rejected update earns no sequence ID") - updates++ - if updates == 1 { - close(firstRequest) - http.Error(w, `{"error_code":"INVALID_PARAMETER_VALUE","message":"first update failed"}`, http.StatusBadRequest) - return - } - assert.JSONEq(t, `{"state":{"name":"second"}}`, body.State) - _, _ = w.Write([]byte(`{"sequence_id":1}`)) - }) - var db dstate.DeploymentState - path := filepath.Join(t.TempDir(), "resources.json") - require.NoError(t, db.Open(ctx, path, dstate.WritingDMS, false, dstate.OpenDmsArgs{})) - db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) - require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", map[string]string{"name": "first"}, nil)) - <-firstRequest - synctest.Wait() - require.ErrorContains(t, db.RecordingError(), "first update failed") - require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", map[string]string{"name": "second"}, nil)) - _, err := db.Finalize(ctx) - require.ErrorContains(t, err, "first update failed") - assert.Equal(t, 2, updates) - }) -} - func TestHistoryErrorsFailOpen(t *testing.T) { for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { t.Run(fmt.Sprint(mode), func(t *testing.T) { diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 2128be2e5ca..74ea60b2c48 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -11,6 +11,10 @@ import ( "github.com/databricks/databricks-sdk-go/service/bundledeployments" ) +// bufferedOperations caps how far ahead of the service a deploy may get; DMS is what the next +// plan reads. +const bufferedOperations = 10 + // BufferOptions controls whether recording continues after an error. type BufferOptions struct { BestEffort bool @@ -23,7 +27,12 @@ type OperationBuffer struct { deploymentID string versionNum int - done chan struct{} + // The buffer. queue holds bundle state keys, and pending the newest update per key, so a + // second write for a resource replaces the first. + queue chan string + done chan struct{} + stopQueue func() + cancel context.CancelFunc bestEffort bool @@ -39,9 +48,9 @@ type OperationBuffer struct { // mu guards the fields below. mu sync.Mutex - // queue is unbounded and contains one key per pending update. - queue []string - draining bool + // File-history recording uses an unbounded queue so it never waits for DMS. + historyQueue []string + draining bool // pending holds the newest update per resource key, absent once the writer takes it. pending map[string]OperationUpdate @@ -62,6 +71,7 @@ func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum i b := &OperationBuffer{ deploymentID: deploymentID, versionNum: versionNum, + queue: make(chan string, bufferedOperations), done: make(chan struct{}), cancel: cancel, bestEffort: options.BestEffort, @@ -69,6 +79,7 @@ func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum i latestState: make(map[string]OperationUpdate), sequenceIDs: make(map[string]int64), } + b.stopQueue = sync.OnceFunc(func() { close(b.queue) }) go b.run(ctx) return b } @@ -98,8 +109,8 @@ func (b *OperationBuffer) RecordFailure(resourceKey, resourceID string, state js // that latest state. latestState and pending are updated under the one lock. func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, keepExistingUpdateIfAny bool, cause error) { b.mu.Lock() - defer b.mu.Unlock() if b.bestEffort && b.err != nil { + b.mu.Unlock() return } if keepExistingUpdateIfAny { @@ -114,13 +125,21 @@ func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, kee update = waiting.Merge(update) } b.pending[resourceKey] = update + if b.bestEffort { + if !queued { + b.historyQueue = append(b.historyQueue, resourceKey) + } + b.mu.Unlock() + return + } + b.mu.Unlock() // when already queued: the writer reads the map when it gets to the key, // so it picks up what was just stored. // // Only when the resource is not already queued, enqueue it. if !queued { - b.queue = append(b.queue, resourceKey) + b.queue <- resourceKey } } @@ -139,31 +158,40 @@ func (b *OperationBuffer) run(ctx context.Context) { defer b.cancel() for { - if err := ctx.Err(); err != nil { - b.setErr(err) - return - } - b.mu.Lock() - if len(b.queue) == 0 { - draining := b.draining - b.mu.Unlock() - if draining { + var resourceKey string + if b.bestEffort { + if err := ctx.Err(); err != nil { + b.setErr(err) return } - // Poll an empty queue without holding the lock or delaying cancellation. - select { - case <-time.After(100 * time.Millisecond): - case <-ctx.Done(): + b.mu.Lock() + if len(b.historyQueue) == 0 { + draining := b.draining + b.mu.Unlock() + if draining { + return + } + // Poll only optional history; DMS-state writes retain immediate delivery. + select { + case <-time.After(100 * time.Millisecond): + case <-ctx.Done(): + } + continue + } + resourceKey = b.historyQueue[0] + b.historyQueue[0] = "" + b.historyQueue = b.historyQueue[1:] + if len(b.historyQueue) == 0 { + b.historyQueue = nil + } + b.mu.Unlock() + } else { + var ok bool + resourceKey, ok = <-b.queue + if !ok { + return } - continue - } - resourceKey := b.queue[0] - b.queue[0] = "" - b.queue = b.queue[1:] - if len(b.queue) == 0 { - b.queue = nil } - b.mu.Unlock() update, ok := b.take(resourceKey) if !ok { // Unreachable: a key is queued only when nothing was waiting for it. Guard so @@ -198,9 +226,15 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio return nil } -// Drain waits for buffered writes, canceling outstanding requests when ctx expires. +// Drain waits for buffered writes. Optional history cancels outstanding requests when ctx expires. // It returns the first recording error and is safe to call twice. func (b *OperationBuffer) Drain(ctx context.Context) error { + if !b.bestEffort { + b.stopQueue() + <-b.done + return b.Err() + } + b.mu.Lock() b.draining = true b.mu.Unlock() diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index 5c8d1c53bf6..29eefccc855 100644 --- a/libs/dms/operation_buffer_test.go +++ b/libs/dms/operation_buffer_test.go @@ -22,6 +22,7 @@ import ( // queued builds a buffer whose queue nothing drains, so a test drives record and take itself. func queued() *OperationBuffer { return &OperationBuffer{ + queue: make(chan string, bufferedOperations), cancel: func() {}, pending: make(map[string]OperationUpdate), latestState: make(map[string]OperationUpdate), @@ -115,15 +116,16 @@ func TestRecordFailureUsesTheGivenStateWhenNoWriteRecorded(t *testing.T) { func TestBufferQueuesWithoutBackpressure(t *testing.T) { synctest.Test(t, func(t *testing.T) { b := queued() + b.bestEffort = true const operationCount = 1000 for i := range operationCount { key := fmt.Sprintf("resources.jobs.job_%d", i) b.RecordOperation(t.Context(), key, false, "id-1", json.RawMessage(`{"state":{}}`)) } assert.NoError(t, b.Err()) - assert.Len(t, b.queue, operationCount) + assert.Len(t, b.historyQueue, operationCount) assert.Len(t, b.pending, operationCount) - for i, key := range b.queue { + for i, key := range b.historyQueue { assert.Equal(t, fmt.Sprintf("resources.jobs.job_%d", i), key) } }) @@ -133,7 +135,7 @@ func TestBufferStopsIdleWriterOnCancellation(t *testing.T) { synctest.Test(t, func(t *testing.T) { ctx, cancel := context.WithCancel(t.Context()) defer cancel() - b := StartOperationBuffer(ctx, "dep-1", 1, BufferOptions{}) + b := StartOperationBuffer(ctx, "dep-1", 1, BufferOptions{BestEffort: true}) synctest.Wait() cancel() <-b.done @@ -159,6 +161,6 @@ func TestBestEffortBufferStopsRecordingAfterError(t *testing.T) { b.RecordOperation(t.Context(), "resources.jobs.bar", false, "id-2", json.RawMessage(`{}`)) b.RecordFailure("resources.jobs.bar", "id-2", nil, errors.New("apply failed")) assert.ErrorContains(t, b.Err(), "exceeds") - assert.Empty(t, b.queue) + assert.Empty(t, b.historyQueue) assert.Empty(t, b.pending) } From 9978aab88f0e25ca71a0bd02a9ae1ce7d9210249 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 11:46:43 +0000 Subject: [PATCH 40/52] Explain releasing the empty history queue backing array --- libs/dms/operation_buffer.go | 1 + 1 file changed, 1 insertion(+) diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 74ea60b2c48..d9fe1e6edd9 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -182,6 +182,7 @@ func (b *OperationBuffer) run(ctx context.Context) { b.historyQueue[0] = "" b.historyQueue = b.historyQueue[1:] if len(b.historyQueue) == 0 { + // Release the backing array so the garbage collector can reclaim it. b.historyQueue = nil } b.mu.Unlock() From 90f22a16a11db72fd063aa17eabfe3def5c66150 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 11:48:57 +0000 Subject: [PATCH 41/52] Clarify history draining polling and cancellation with short comments --- bundle/direct/dstate/state.go | 4 +++- libs/dms/operation_buffer.go | 20 ++++++++++++++++---- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 759d3f40eb4..ccec75c872e 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -189,6 +189,7 @@ func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploy BestEffort: db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced, } db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options) + // Use the full plan count; the queue length changes as uploads finish. db.historyDrainTimeout = time.Duration(operationCount) * historyDrainTimeoutPerOperation db.DeploymentID = deploymentID db.VersionID = versionID @@ -889,7 +890,8 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported } case WritingFileWithHistory, WritingFileWithHistoryEnforced: if db.operationBuffer != nil { - // File state is persisted before waiting on optional history recording. + // File state is already persisted, so history errors only warn. + // Bound the wait so a stalled service cannot hold up deployment. drainCtx, cancel := context.WithTimeout(ctx, db.historyDrainTimeout) drainErr := db.operationBuffer.Drain(drainCtx) cancel() diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index d9fe1e6edd9..85ca9a50118 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -33,6 +33,7 @@ type OperationBuffer struct { done chan struct{} stopQueue func() + // cancel stops uploads when optional history fails or its drain deadline expires. cancel context.CancelFunc bestEffort bool @@ -50,7 +51,9 @@ type OperationBuffer struct { // File-history recording uses an unbounded queue so it never waits for DMS. historyQueue []string - draining bool + + // draining lets the worker exit once the history queue is empty. + draining bool // pending holds the newest update per resource key, absent once the writer takes it. pending map[string]OperationUpdate @@ -67,6 +70,7 @@ type OperationBuffer struct { // StartOperationBuffer opens the buffer for the version the caller just created. The version // must already exist: operations record under it, and nothing here creates it. func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum int, options BufferOptions) *OperationBuffer { + // History can stop independently without canceling the deployment. ctx, cancel := context.WithCancel(ctx) b := &OperationBuffer{ deploymentID: deploymentID, @@ -109,6 +113,7 @@ func (b *OperationBuffer) RecordFailure(resourceKey, resourceID string, state js // that latest state. latestState and pending are updated under the one lock. func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, keepExistingUpdateIfAny bool, cause error) { b.mu.Lock() + // Stop building a history backlog after the first recording failure. if b.bestEffort && b.err != nil { b.mu.Unlock() return @@ -160,6 +165,7 @@ func (b *OperationBuffer) run(ctx context.Context) { for { var resourceKey string if b.bestEffort { + // Cancellation comes from the deployment, a drain timeout, or a history error. if err := ctx.Err(); err != nil { b.setErr(err) return @@ -171,7 +177,8 @@ func (b *OperationBuffer) run(ctx context.Context) { if draining { return } - // Poll only optional history; DMS-state writes retain immediate delivery. + // Poll while idle so producers need no wake-up signaling. + // Cancellation ends this wait immediately. select { case <-time.After(100 * time.Millisecond): case <-ctx.Done(): @@ -227,8 +234,11 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio return nil } -// Drain waits for buffered writes. Optional history cancels outstanding requests when ctx expires. -// It returns the first recording error and is safe to call twice. +// Drain waits for buffered uploads before version completion. +// - DMS state: close the channel and wait for queued writes. +// - File history: finish the queue, or cancel uploads when ctx expires. +// +// It returns the first recording error and is safe to call more than once. func (b *OperationBuffer) Drain(ctx context.Context) error { if !b.bestEffort { b.stopQueue() @@ -243,6 +253,7 @@ func (b *OperationBuffer) Drain(ctx context.Context) error { case <-b.done: return b.Err() case <-ctx.Done(): + // Stop the in-flight request so history cannot outlive the drain deadline. b.setErr(ctx.Err()) b.cancel() return b.Err() @@ -257,6 +268,7 @@ func (b *OperationBuffer) setErr(err error) { if b.err == nil { b.err = err if b.bestEffort { + // Stop polling and uploads even when the producer reports the error. b.cancel() } } From 203338b776d98e89c3d6ce8c834be404e70aacf2 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 11:50:13 +0000 Subject: [PATCH 42/52] Clarify that file history does not backpressure the deployment engine --- libs/dms/operation_buffer.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 85ca9a50118..361e9f753d7 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -49,7 +49,8 @@ type OperationBuffer struct { // mu guards the fields below. mu sync.Mutex - // File-history recording uses an unbounded queue so it never waits for DMS. + // File-history recording uses an unbounded queue so DMS applies no backpressure + // to the deployment engine. historyQueue []string // draining lets the worker exit once the history queue is empty. From 2a4e9cb37306901850d342ccce0ec3bc5aabaa62 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 11:51:19 +0000 Subject: [PATCH 43/52] Explain the expected cancellation error in the history worker --- libs/dms/operation_buffer.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 361e9f753d7..31a14cd1682 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -166,7 +166,9 @@ func (b *OperationBuffer) run(ctx context.Context) { for { var resourceKey string if b.bestEffort { - // Cancellation comes from the deployment, a drain timeout, or a history error. + // Calling b.cancel() makes ctx.Err() return an error. This is the signal to stop. + // We expect this after a drain timeout or history recording error. + // Cancellation or expiry of the deployment's context also stops the worker here. if err := ctx.Err(); err != nil { b.setErr(err) return From cc302bc9fe42954945cf052420cda9014f00b2e8 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 11:51:55 +0000 Subject: [PATCH 44/52] Document polling as the initial history queue design --- libs/dms/operation_buffer.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 31a14cd1682..0138a8436fe 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -180,7 +180,8 @@ func (b *OperationBuffer) run(ctx context.Context) { if draining { return } - // Poll while idle so producers need no wake-up signaling. + // Start with polling to keep the design simple. + // We can optimize this later by waking the worker when work is queued. // Cancellation ends this wait immediately. select { case <-time.After(100 * time.Millisecond): From 3066b3ddcab4b924e265b16f9df1fc64b063bf0f Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 12:03:53 +0000 Subject: [PATCH 45/52] Continue recording after upload failures and reduce mode test scope --- bundle/direct/dstate/mode_test.go | 362 ++++-------------------------- bundle/direct/dstate/state.go | 11 +- bundle/phases/deploy.go | 1 - bundle/phases/destroy.go | 4 +- bundle/phases/telemetry.go | 7 +- libs/dms/operation_buffer.go | 29 +-- libs/dms/operation_buffer_test.go | 18 +- 7 files changed, 59 insertions(+), 373 deletions(-) diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go index 670adeb043a..fd22981d24c 100644 --- a/bundle/direct/dstate/mode_test.go +++ b/bundle/direct/dstate/mode_test.go @@ -10,7 +10,6 @@ import ( "os" "path/filepath" "strconv" - "sync" "testing" "testing/synctest" "time" @@ -68,7 +67,7 @@ func TestGetSerialWithDivergingCounters(t *testing.T) { {dstate.WritingFileWithHistoryEnforced, 7, 7}, } { t.Run(fmt.Sprint(tc.mode), func(t *testing.T) { - ctx := modeContext(t, nil) + ctx := t.Context() var db dstate.DeploymentState path := filepath.Join(t.TempDir(), "resources.json") require.NoError(t, db.Open(ctx, path, tc.mode, false, dstate.OpenDmsArgs{})) @@ -82,68 +81,23 @@ func TestGetSerialWithDivergingCounters(t *testing.T) { } } -func TestModeAssertions(t *testing.T) { - for _, tc := range []struct { - mode dstate.Mode - writable bool - dms bool - history bool - }{ - {dstate.ReadingFile, false, false, false}, - {dstate.WritingFile, true, false, false}, - {dstate.WritingDMS, true, true, true}, - {dstate.WritingFileWithHistory, true, false, true}, - {dstate.WritingFileWithHistoryEnforced, true, false, true}, - } { - t.Run(fmt.Sprint(tc.mode), func(t *testing.T) { - ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { - t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) - }) - path := filepath.Join(t.TempDir(), "resources.json") - var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, tc.mode, false, dstate.OpenDmsArgs{})) - assert.Equal(t, tc.dms, db.IsDeploymentMetadataService()) - assert.Equal(t, tc.history, db.RecordsHistory()) - assert.NotPanics(t, db.AssertOpenedForReadOrWrite) - if tc.writable { - assert.NotPanics(t, db.AssertOpenedForWrite) - if tc.dms { - assert.NotPanics(t, db.AssertOpenedForRead) - } else { - assert.Panics(t, db.AssertOpenedForRead) - } - } else { - assert.NotPanics(t, db.AssertOpenedForRead) - assert.Panics(t, db.AssertOpenedForWrite) - } - if tc.writable && !tc.dms { - assert.FileExists(t, path+".wal") - } else { - assert.NoFileExists(t, path+".wal") - } - _, err := db.Finalize(ctx) - require.NoError(t, err) - assert.Panics(t, db.AssertOpenedForReadOrWrite) - assert.NoFileExists(t, path+".wal") - }) - } -} - func TestFileModesPersistStateAndHistory(t *testing.T) { for _, mode := range []dstate.Mode{dstate.WritingFile, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { t.Run(fmt.Sprint(mode), func(t *testing.T) { - var mu sync.Mutex updates := make(map[string]map[string]any) ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost { + assert.Equal(t, "/api/2.0/bundle/deployments/dep-1/versions/1/complete", r.URL.Path) + _, _ = w.Write([]byte(`{}`)) + return + } assert.Equal(t, http.MethodPatch, r.Method) var body map[string]any if !assert.NoError(t, json.NewDecoder(r.Body).Decode(&body)) { w.WriteHeader(http.StatusBadRequest) return } - mu.Lock() updates[filepath.Base(r.URL.Path)] = body - mu.Unlock() w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"sequence_id":1}`)) }) @@ -162,6 +116,12 @@ func TestFileModesPersistStateAndHistory(t *testing.T) { assert.Len(t, exported, 1) assert.Equal(t, "keep-id", exported["resources.jobs.kept"].ID) assert.NoFileExists(t, path+".wal") + completed, err := db.CompleteVersion(ctx, true) + require.NoError(t, err) + assert.Equal(t, mode != dstate.WritingFile, completed) + completed, err = db.CompleteVersion(ctx, true) + require.NoError(t, err) + assert.False(t, completed) var reopened dstate.DeploymentState require.NoError(t, reopened.Open(t.Context(), path, dstate.ReadingFile, false, dstate.OpenDmsArgs{})) @@ -173,8 +133,6 @@ func TestFileModesPersistStateAndHistory(t *testing.T) { _, err = reopened.Finalize(t.Context()) require.NoError(t, err) - mu.Lock() - defer mu.Unlock() if mode == dstate.WritingFile { assert.Empty(t, updates) } else { @@ -186,87 +144,16 @@ func TestFileModesPersistStateAndHistory(t *testing.T) { } } -func TestDMSReadsAndWritesWithoutWAL(t *testing.T) { - ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - if r.Method == http.MethodGet { - _, _ = w.Write([]byte(`{"resources":[{"resource_key":"jobs.remote","resource_id":"remote-id","state":"{\"state\":{\"name\":\"remote\"}}"}]}`)) - } else { - _, _ = w.Write([]byte(`{"sequence_id":1}`)) - } - }) - path := filepath.Join(t.TempDir(), "resources.json") - var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, dstate.WritingDMS, false, dstate.OpenDmsArgs{DeploymentID: "dep-1", LastVersionID: 7})) - assert.Equal(t, "remote-id", db.GetResourceID("resources.jobs.remote")) - assert.Equal(t, 7, db.GetSerial()) - require.NoError(t, db.UpgradeToWrite()) - assert.NoFileExists(t, path+".wal") - db.InitializeOperationBuffer(ctx, "dep-1", 8, 2) - require.NoError(t, db.SaveState(ctx, "resources.jobs.new", "new-id", map[string]string{"name": "new"}, nil)) - require.NoError(t, db.DeleteState(ctx, "resources.jobs.remote", false)) - exported, err := db.Finalize(ctx) - require.NoError(t, err) - assert.Len(t, exported, 1) - assert.Equal(t, "new-id", exported["resources.jobs.new"].ID) - assert.NoFileExists(t, path+".wal") - data, err := os.ReadFile(path) - require.NoError(t, err) - var marker dstate.Database - require.NoError(t, json.Unmarshal(data, &marker)) - assert.Empty(t, marker.State) - assert.Equal(t, 8, marker.Serial) - assert.Contains(t, marker.Features, dstate.FeatureDeploymentHistory) -} - -func TestUpgradeToWrite(t *testing.T) { - for _, tc := range []struct { - name string - mode dstate.Mode - bestEffort bool - wantErr string - }{ - {"file", dstate.ReadingFile, false, ""}, - {"DMS", dstate.WritingDMS, false, ""}, - {"history", dstate.WritingFileWithHistory, false, ""}, - {"enforced history", dstate.WritingFileWithHistoryEnforced, false, ""}, - {"file already writable", dstate.WritingFile, false, "already open for write"}, - {"enforcement retained", dstate.WritingFileWithHistory, true, ""}, - {"file upgrade with history marker", dstate.ReadingFile, true, ""}, - {"file open with history marker", dstate.WritingFile, true, "already open for write"}, - } { - t.Run(tc.name, func(t *testing.T) { - ctx := modeContext(t, nil) - path := filepath.Join(t.TempDir(), "resources.json") - if tc.bestEffort { - require.NoError(t, os.WriteFile(path, []byte(`{"state_version":3,"features":{"deployment_history_best_effort":{}},"state":{}}`), 0o600)) - } - var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, tc.mode, false, dstate.OpenDmsArgs{})) - err := db.UpgradeToWrite() - if tc.wantErr != "" { - require.ErrorContains(t, err, tc.wantErr) - } else { - require.NoError(t, err) - assert.NotPanics(t, db.AssertOpenedForWrite) - } - _, err = db.Finalize(ctx) - require.NoError(t, err) - if tc.bestEffort { - data, err := os.ReadFile(path) - require.NoError(t, err) - var persisted dstate.Database - require.NoError(t, json.Unmarshal(data, &persisted)) - assert.Contains(t, persisted.Features, dstate.FeatureDeploymentHistoryBestEffort) - } - }) - } -} - func TestHistoryErrorsFailOpen(t *testing.T) { for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { t.Run(fmt.Sprint(mode), func(t *testing.T) { + goodUploads := 0 ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPatch && filepath.Base(r.URL.Path) == "jobs.good" { + goodUploads++ + _, _ = w.Write([]byte(`{"sequence_id":1}`)) + return + } w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusBadRequest) _, _ = w.Write([]byte(`{"error_code":"INVALID_PARAMETER_VALUE","message":"history unavailable"}`)) @@ -274,9 +161,11 @@ func TestHistoryErrorsFailOpen(t *testing.T) { path := filepath.Join(t.TempDir(), "resources.json") var db dstate.DeploymentState require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) - db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) + db.InitializeOperationBuffer(ctx, "dep-1", 1, 2) require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", map[string]string{"name": "job"}, nil)) + require.NoError(t, db.SaveState(ctx, "resources.jobs.good", "id-2", struct{}{}, nil)) _, err := db.Finalize(ctx) + assert.Equal(t, 1, goodUploads, "a bad upload must not stop the next resource") if mode == dstate.WritingDMS { require.ErrorContains(t, err, "history unavailable") assert.Error(t, db.RecordingError()) @@ -302,201 +191,26 @@ func TestHistoryErrorsFailOpen(t *testing.T) { } } -func TestHistoryDrainTimeoutPreservesFileState(t *testing.T) { - for _, mode := range []dstate.Mode{dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { - for _, tc := range []struct { - name string - operationCount int - callerTimeout time.Duration - wantDuration time.Duration - }{ - {name: "one operation", operationCount: 1, wantDuration: 5 * time.Second}, - {name: "all staged operations", operationCount: 4, wantDuration: 20 * time.Second}, - {name: "shorter caller deadline", operationCount: 4, callerTimeout: 20 * time.Millisecond, wantDuration: 20 * time.Millisecond}, - } { - t.Run(fmt.Sprintf("%d/%s", mode, tc.name), func(t *testing.T) { - synctest.Test(t, func(t *testing.T) { - started := make(chan struct{}) - finished := make(chan struct{}) - ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { - _, _ = io.Copy(io.Discard, r.Body) - close(started) - <-r.Context().Done() - close(finished) - }) - path := filepath.Join(t.TempDir(), "resources.json") - var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) - db.InitializeOperationBuffer(ctx, "dep-1", 1, tc.operationCount) - // One buffered update must still get the budget for all staged operations. - require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) - <-started - drainCtx := ctx - if tc.callerTimeout > 0 { - var cancel context.CancelFunc - drainCtx, cancel = context.WithTimeout(ctx, tc.callerTimeout) - defer cancel() - } - start := time.Now() - _, err := db.Finalize(drainCtx) - require.NoError(t, err) - assert.Equal(t, tc.wantDuration, time.Since(start)) - assert.FileExists(t, path) - assert.NoFileExists(t, path+".wal") - <-finished - }) - }) - } - } -} - -func TestCompleteVersionHonorsCallerDeadline(t *testing.T) { - for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { - t.Run(fmt.Sprint(mode), func(t *testing.T) { - synctest.Test(t, func(t *testing.T) { - ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodPost { - _, _ = io.Copy(io.Discard, r.Body) - <-r.Context().Done() - } else { - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"sequence_id":1}`)) - } - }) - var db dstate.DeploymentState - path := filepath.Join(t.TempDir(), "resources.json") - require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) - db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) - require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) - _, err := db.Finalize(ctx) - require.NoError(t, err) - completionCtx, cancel := context.WithTimeout(ctx, 90*time.Second) - defer cancel() - start := time.Now() - completed, err := db.CompleteVersion(completionCtx, true) - assert.Equal(t, 90*time.Second, time.Since(start)) - assert.False(t, completed) - assert.False(t, db.RecordsHistory()) - assert.Empty(t, db.DeploymentID) - assert.Zero(t, db.VersionID) - if mode == dstate.WritingDMS { - require.ErrorIs(t, err, context.DeadlineExceeded) - } else { - require.NoError(t, err) - } - }) - }) - } -} - func TestEnforcedHistorySurvivesWALRecovery(t *testing.T) { - for _, withResource := range []bool{false, true} { - t.Run(map[bool]string{false: "feature only", true: "with resource"}[withResource], func(t *testing.T) { - ctx := modeContext(t, nil) - path := filepath.Join(t.TempDir(), "resources.json") - var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, dstate.WritingFileWithHistoryEnforced, false, dstate.OpenDmsArgs{})) - if withResource { - require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) - } - wal, err := os.ReadFile(path + ".wal") - require.NoError(t, err) - _, err = db.Finalize(ctx) - require.NoError(t, err) - // Reproduce an interrupted deployment using only its WAL. - recoveredPath := filepath.Join(t.TempDir(), "resources.json") - require.NoError(t, os.WriteFile(recoveredPath+".wal", wal, 0o600)) - var recovered dstate.DeploymentState - require.NoError(t, recovered.Open(t.Context(), recoveredPath, dstate.ReadingFile, true, dstate.OpenDmsArgs{})) - assert.Contains(t, recovered.StateFeatures(), dstate.FeatureDeploymentHistoryBestEffort) - assert.Equal(t, 1, recovered.GetSerial()) - assert.FileExists(t, recoveredPath) - assert.NoFileExists(t, recoveredPath+".wal") - require.NoError(t, recovered.UpgradeToWrite()) - if withResource { - assert.Equal(t, "id-1", recovered.GetResourceID("resources.jobs.job")) - } - _, err = recovered.Finalize(ctx) - require.NoError(t, err) - }) - } -} - -func TestOpenRejectsIncompatibleModesAndFeatures(t *testing.T) { - for _, tc := range []struct { - name string - mode dstate.Mode - features string - wantErr string - }{ - {"unknown feature", dstate.ReadingFile, `{"future":{}}`, "requires features"}, - {"file reading DMS with best effort marker", dstate.ReadingFile, `{"deployment_history":{},"deployment_history_best_effort":{}}`, "not supported"}, - {"file writing DMS", dstate.WritingFile, `{"deployment_history":{}}`, "not supported"}, - } { - t.Run(tc.name, func(t *testing.T) { - path := filepath.Join(t.TempDir(), "resources.json") - seed := `{"state_version":3,"features":` + tc.features + `,"state":{}}` - require.NoError(t, os.WriteFile(path, []byte(seed), 0o600)) - var db dstate.DeploymentState - err := db.Open(t.Context(), path, tc.mode, false, dstate.OpenDmsArgs{}) - require.ErrorContains(t, err, tc.wantErr) - assert.False(t, db.IsOpen()) - assert.NoFileExists(t, path+".wal") - }) - } -} - -func TestCompleteVersionAfterFinalize(t *testing.T) { - for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { - t.Run(fmt.Sprint(mode), func(t *testing.T) { - completions := 0 - ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "application/json") - if r.Method == http.MethodPost { - completions++ - assert.Equal(t, "/api/2.0/bundle/deployments/dep-1/versions/1/complete", r.URL.Path) - var body map[string]string - assert.NoError(t, json.NewDecoder(r.Body).Decode(&body)) - assert.Equal(t, "VERSION_COMPLETE_SUCCESS", body["completion_reason"]) - _, _ = w.Write([]byte(`{}`)) - } else { - _, _ = w.Write([]byte(`{"sequence_id":1}`)) - } - }) - var db dstate.DeploymentState - for cycle := range 2 { - path := filepath.Join(t.TempDir(), "resources.json") - require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) - db.InitializeOperationBuffer(ctx, "dep-1", 1, 1) - require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", struct{}{}, nil)) - _, err := db.Finalize(ctx) - require.NoError(t, err) - assert.False(t, db.IsOpen()) - completed, err := db.CompleteVersion(ctx, true) - require.NoError(t, err) - assert.True(t, completed) - completed, err = db.CompleteVersion(ctx, true) - require.NoError(t, err) - assert.False(t, completed) - assert.Equal(t, cycle+1, completions) - assert.False(t, db.RecordsHistory()) - assert.Empty(t, db.DeploymentID) - assert.Zero(t, db.VersionID) - } - }) - } -} - -func TestFailedUpgradeLeavesReadModeAvailable(t *testing.T) { + ctx := t.Context() path := filepath.Join(t.TempDir(), "resources.json") var db dstate.DeploymentState - require.NoError(t, db.Open(t.Context(), path, dstate.ReadingFile, true, dstate.OpenDmsArgs{})) - require.NoError(t, os.WriteFile(path+".wal", []byte("occupied"), 0o600)) - require.Error(t, db.UpgradeToWrite()) - assert.NotPanics(t, db.AssertOpenedForRead) - require.NoError(t, os.Remove(path+".wal")) - require.NoError(t, db.UpgradeToWrite()) - _, err := db.Finalize(t.Context()) + require.NoError(t, db.Open(ctx, path, dstate.WritingFileWithHistoryEnforced, false, dstate.OpenDmsArgs{})) + wal, err := os.ReadFile(path + ".wal") + require.NoError(t, err) + _, err = db.Finalize(ctx) + require.NoError(t, err) + + // Recover an interrupted first deployment: its WAL contains only the feature marker. + recoveredPath := filepath.Join(t.TempDir(), "resources.json") + require.NoError(t, os.WriteFile(recoveredPath+".wal", wal, 0o600)) + var recovered dstate.DeploymentState + require.NoError(t, recovered.Open(ctx, recoveredPath, dstate.ReadingFile, true, dstate.OpenDmsArgs{})) + assert.Contains(t, recovered.StateFeatures(), dstate.FeatureDeploymentHistoryBestEffort) + assert.Equal(t, 1, recovered.GetSerial()) + assert.FileExists(t, recoveredPath) + assert.NoFileExists(t, recoveredPath+".wal") + _, err = recovered.Finalize(ctx) require.NoError(t, err) } diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index ccec75c872e..dc99ddc2ade 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -564,14 +564,11 @@ func (db *DeploymentState) reset() { db.Data = Database{} db.stateIDs = nil db.historyDrainTimeout = 0 - // Keep the buffer, mode, and IDs after Finalize so CompleteVersion can use them. - // Clear them if no version was started, or CompleteVersion has copied what it needs. - if db.operationBuffer == nil || db.versionCompleted { - db.mode = ReadingFile + // Telemetry and destroy still need the mode and IDs after version completion. + // Clear the recording buffer once CompleteVersion has copied it. + if db.versionCompleted { db.operationBuffer = nil db.versionCompleted = false - db.DeploymentID = "" - db.VersionID = 0 } } @@ -1000,6 +997,8 @@ func (db *DeploymentState) AssertOpenedForReadOrWrite() { func (db *DeploymentState) AssertOpenedForRead() { db.AssertOpenedForReadOrWrite() + // DMS has no separate read-only mode. + // Planning reads DMS state while it is already in WritingDMS. if db.mode.isWritable() && db.mode != WritingDMS { panic("internal error: DeploymentState must be opened in read mode") } diff --git a/bundle/phases/deploy.go b/bundle/phases/deploy.go index dca864ae670..13c5121f569 100644 --- a/bundle/phases/deploy.go +++ b/bundle/phases/deploy.go @@ -155,7 +155,6 @@ func uploadLibraries(ctx context.Context, b *bundle.Bundle, libs map[string][]li // If readPlanPath is provided, the plan is loaded from that file instead of being calculated. func Deploy(ctx context.Context, b *bundle.Bundle, outputHandler sync.OutputHandler, libs map[string][]libraries.LocationToUpdate, plan *deployplan.Plan, dmsDeployment *bundledeployments.Deployment) { log.Info(ctx, "Phase: deploy") - b.Metrics.SetBoolValue(metrics.DeploymentHistoryEnabled, b.DeploymentBundle.StateDB.IsDeploymentMetadataService()) // Core mutators that CRUD resources and modify deployment state. These // mutators need informed consent if they are potentially destructive. The diff --git a/bundle/phases/destroy.go b/bundle/phases/destroy.go index 408b4fd3191..07bbbbca385 100644 --- a/bundle/phases/destroy.go +++ b/bundle/phases/destroy.go @@ -158,7 +158,6 @@ func destroyCore(ctx context.Context, b *bundle.Bundle, plan *deployplan.Plan) { if b.DeploymentBundle.StateDB.IsDeploymentMetadataService() { // Complete version before deleting remote files; the deployment node is under statePath. - deploymentID := b.DeploymentBundle.StateDB.DeploymentID completed, err := b.DeploymentBundle.StateDB.CompleteVersion(ctx, true) if err != nil { logdiag.LogError(ctx, err) @@ -166,6 +165,7 @@ func destroyCore(ctx context.Context, b *bundle.Bundle, plan *deployplan.Plan) { } // A completed destroy's resources are gone, so its deployment record is deleted too. if completed { + deploymentID := b.DeploymentBundle.StateDB.DeploymentID w := b.WorkspaceClient(ctx) if err := deleteDeployment(ctx, w, deploymentID); err != nil { logdiag.LogError(ctx, fmt.Errorf("failed to delete deployment: %w", err)) @@ -294,11 +294,11 @@ func Destroy(ctx context.Context, b *bundle.Bundle) { // destroyCore has completed the version. defer func() { if b.DeploymentBundle.StateDB.IsDeploymentMetadataService() { - deploymentID := b.DeploymentBundle.StateDB.DeploymentID completed, err := b.DeploymentBundle.StateDB.CompleteVersion(ctx, !logdiag.HasError(ctx)) if err != nil { logdiag.LogError(ctx, err) } else if completed { + deploymentID := b.DeploymentBundle.StateDB.DeploymentID w := b.WorkspaceClient(ctx) if err := deleteDeployment(ctx, w, deploymentID); err != nil { logdiag.LogError(ctx, fmt.Errorf("failed to delete deployment: %w", err)) diff --git a/bundle/phases/telemetry.go b/bundle/phases/telemetry.go index 013512b8aba..99464f3ef24 100644 --- a/bundle/phases/telemetry.go +++ b/bundle/phases/telemetry.go @@ -198,12 +198,7 @@ func aiRuntimeTaskMetrics(jobs map[string]*resources.Job) (present, scheduled, m // LogDeployTelemetry logs a telemetry event for a bundle deploy command. func LogDeployTelemetry(ctx context.Context, b *bundle.Bundle, errMsg string) { errMsg = telemetry.ScrubErrorMessage(errMsg) - // Deploy captures this before state cleanup. Fall back to the current state for earlier failures. - if !slices.ContainsFunc(b.Metrics.BoolValues, func(v protos.BoolMapEntry) bool { - return v.Key == metrics.DeploymentHistoryEnabled - }) { - b.Metrics.SetBoolValue(metrics.DeploymentHistoryEnabled, b.DeploymentBundle.StateDB.IsDeploymentMetadataService()) - } + b.Metrics.SetBoolValue(metrics.DeploymentHistoryEnabled, b.DeploymentBundle.StateDB.IsDeploymentMetadataService()) resourcesCount := int64(0) _, err := dyn.MapByPattern(b.Config.Value(), dyn.NewPattern(dyn.Key("resources"), dyn.AnyKey(), dyn.AnyKey()), func(p dyn.Path, v dyn.Value) (dyn.Value, error) { diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 0138a8436fe..8e8973ecf7d 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -15,7 +15,7 @@ import ( // plan reads. const bufferedOperations = 10 -// BufferOptions controls whether recording continues after an error. +// BufferOptions selects optional file-history recording behavior. type BufferOptions struct { BestEffort bool } @@ -33,7 +33,7 @@ type OperationBuffer struct { done chan struct{} stopQueue func() - // cancel stops uploads when optional history fails or its drain deadline expires. + // cancel stops in-flight history uploads when the drain deadline expires. cancel context.CancelFunc bestEffort bool @@ -90,7 +90,7 @@ func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum i } // RecordOperation records a state write. state is the serialized envelope, nil for a delete. -// Required recording continues after an earlier failure; optional history stops on the first error. +// An earlier failure does not stop it: keep recording, best effort. func (b *OperationBuffer) RecordOperation(ctx context.Context, resourceKey string, inProgress bool, resourceID string, state json.RawMessage) { update, err := NewStateUpdate(resourceID, state, inProgress) if err != nil { @@ -114,11 +114,6 @@ func (b *OperationBuffer) RecordFailure(resourceKey, resourceID string, state js // that latest state. latestState and pending are updated under the one lock. func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, keepExistingUpdateIfAny bool, cause error) { b.mu.Lock() - // Stop building a history backlog after the first recording failure. - if b.bestEffort && b.err != nil { - b.mu.Unlock() - return - } if keepExistingUpdateIfAny { if latest, recorded := b.latestState[resourceKey]; recorded { update = latest.Merge(failureUpdate(cause)) @@ -167,7 +162,7 @@ func (b *OperationBuffer) run(ctx context.Context) { var resourceKey string if b.bestEffort { // Calling b.cancel() makes ctx.Err() return an error. This is the signal to stop. - // We expect this after a drain timeout or history recording error. + // We expect this after a drain timeout. // Cancellation or expiry of the deployment's context also stops the worker here. if err := ctx.Err(); err != nil { b.setErr(err) @@ -211,12 +206,9 @@ func (b *OperationBuffer) run(ctx context.Context) { continue } - // Required recording keeps going so one bad write does not drop everything behind it. + // Keep going after a failure, so one bad write does not drop everything behind it. if err := b.write(ctx, resourceKey, update); err != nil { b.setErr(fmt.Errorf("recording operation for %s: %w", resourceKey, err)) - if b.bestEffort { - return - } } } } @@ -245,36 +237,35 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio // It returns the first recording error and is safe to call more than once. func (b *OperationBuffer) Drain(ctx context.Context) error { if !b.bestEffort { + // DMS is authoritative: wait for all queued writes and return any recording error. b.stopQueue() <-b.done return b.Err() } + // File state is already saved. Try the remaining history writes until ctx expires. b.mu.Lock() b.draining = true b.mu.Unlock() select { case <-b.done: + // Report the first upload error even if later uploads succeeded. return b.Err() case <-ctx.Done(): - // Stop the in-flight request so history cannot outlive the drain deadline. + // Deadline or cancellation: stop the upload so deployment can finish promptly. b.setErr(ctx.Err()) b.cancel() return b.Err() } } -// setErr keeps the first error; one failure is enough to fail the deploy. +// setErr keeps the first recording error for Drain to report. func (b *OperationBuffer) setErr(err error) { b.mu.Lock() defer b.mu.Unlock() if b.err == nil { b.err = err - if b.bestEffort { - // Stop polling and uploads even when the producer reports the error. - b.cancel() - } } } diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index 29eefccc855..14d17aa303b 100644 --- a/libs/dms/operation_buffer_test.go +++ b/libs/dms/operation_buffer_test.go @@ -23,7 +23,6 @@ import ( func queued() *OperationBuffer { return &OperationBuffer{ queue: make(chan string, bufferedOperations), - cancel: func() {}, pending: make(map[string]OperationUpdate), latestState: make(map[string]OperationUpdate), } @@ -143,24 +142,13 @@ func TestBufferStopsIdleWriterOnCancellation(t *testing.T) { }) } -func TestBestEffortBufferStopsIdleWriterAfterError(t *testing.T) { - synctest.Test(t, func(t *testing.T) { - b := StartOperationBuffer(t.Context(), "dep-1", 1, BufferOptions{BestEffort: true}) - synctest.Wait() - cause := errors.New("recording failed") - b.setErr(cause) - <-b.done - assert.ErrorIs(t, b.Drain(t.Context()), cause) - }) -} - -func TestBestEffortBufferStopsRecordingAfterError(t *testing.T) { +func TestBestEffortBufferKeepsRecordingAfterError(t *testing.T) { b := queued() b.bestEffort = true b.RecordOperation(t.Context(), "resources.jobs.foo", false, "id-1", json.RawMessage(strings.Repeat("x", maxStateSize+1))) b.RecordOperation(t.Context(), "resources.jobs.bar", false, "id-2", json.RawMessage(`{}`)) b.RecordFailure("resources.jobs.bar", "id-2", nil, errors.New("apply failed")) assert.ErrorContains(t, b.Err(), "exceeds") - assert.Empty(t, b.historyQueue) - assert.Empty(t, b.pending) + assert.Len(t, b.historyQueue, 1) + assert.Contains(t, b.pending, "resources.jobs.bar") } From 7afb46e496efc5577a4241cc0fb8d15feb879511 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 12:10:28 +0000 Subject: [PATCH 46/52] Remove broad mode tests and retain two small state edge cases --- bundle/direct/dstate/mode_test.go | 259 ----------------------------- bundle/direct/dstate/state_test.go | 33 ++++ 2 files changed, 33 insertions(+), 259 deletions(-) delete mode 100644 bundle/direct/dstate/mode_test.go diff --git a/bundle/direct/dstate/mode_test.go b/bundle/direct/dstate/mode_test.go deleted file mode 100644 index fd22981d24c..00000000000 --- a/bundle/direct/dstate/mode_test.go +++ /dev/null @@ -1,259 +0,0 @@ -package dstate_test - -import ( - "context" - "encoding/json" - "fmt" - "io" - "net/http" - "net/http/httptest" - "os" - "path/filepath" - "strconv" - "testing" - "testing/synctest" - "time" - - "github.com/databricks/cli/bundle/deployplan" - "github.com/databricks/cli/bundle/direct/dstate" - "github.com/databricks/cli/libs/cmdctx" - "github.com/databricks/databricks-sdk-go" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -type modeTransport struct { - handler http.HandlerFunc -} - -func (m modeTransport) RoundTrip(r *http.Request) (*http.Response, error) { - w := httptest.NewRecorder() - if r.URL.Path == "/.well-known/databricks-config" || m.handler == nil { - w.WriteHeader(http.StatusNotFound) - } else { - m.handler.ServeHTTP(w, r) - } - if err := r.Context().Err(); err != nil { - return nil, err - } - response := w.Result() - response.Request = r - return response, nil -} - -// modeContext supplies an in-memory DMS endpoint without resolving workspace authentication. -func modeContext(t *testing.T, handler http.HandlerFunc) context.Context { - t.Helper() - client, err := databricks.NewWorkspaceClient(&databricks.Config{ - Host: "https://state.test", Token: "test-token", AuthType: "pat", - HTTPTransport: modeTransport{handler: handler}, - // Keep SDK inactivity timeouts outside the drain budgets exercised here. - HTTPTimeoutSeconds: 600, - }) - require.NoError(t, err) - return cmdctx.SetWorkspaceClient(t.Context(), client) -} - -func TestGetSerialWithDivergingCounters(t *testing.T) { - for _, tc := range []struct { - mode dstate.Mode - versionAhead int - serialAhead int - }{ - {dstate.ReadingFile, 4, 7}, - {dstate.WritingFile, 4, 7}, - {dstate.WritingDMS, 7, 4}, - {dstate.WritingFileWithHistory, 7, 7}, - {dstate.WritingFileWithHistoryEnforced, 7, 7}, - } { - t.Run(fmt.Sprint(tc.mode), func(t *testing.T) { - ctx := t.Context() - var db dstate.DeploymentState - path := filepath.Join(t.TempDir(), "resources.json") - require.NoError(t, db.Open(ctx, path, tc.mode, false, dstate.OpenDmsArgs{})) - db.Data.Serial, db.VersionID = 4, 7 - assert.Equal(t, tc.versionAhead, db.GetSerial(), "DMS version ahead of file serial") - db.Data.Serial, db.VersionID = 7, 4 - assert.Equal(t, tc.serialAhead, db.GetSerial(), "file serial ahead of DMS version") - _, err := db.Finalize(ctx) - require.NoError(t, err) - }) - } -} - -func TestFileModesPersistStateAndHistory(t *testing.T) { - for _, mode := range []dstate.Mode{dstate.WritingFile, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { - t.Run(fmt.Sprint(mode), func(t *testing.T) { - updates := make(map[string]map[string]any) - ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodPost { - assert.Equal(t, "/api/2.0/bundle/deployments/dep-1/versions/1/complete", r.URL.Path) - _, _ = w.Write([]byte(`{}`)) - return - } - assert.Equal(t, http.MethodPatch, r.Method) - var body map[string]any - if !assert.NoError(t, json.NewDecoder(r.Body).Decode(&body)) { - w.WriteHeader(http.StatusBadRequest) - return - } - updates[filepath.Base(r.URL.Path)] = body - w.Header().Set("Content-Type", "application/json") - _, _ = w.Write([]byte(`{"sequence_id":1}`)) - }) - path := filepath.Join(t.TempDir(), "resources.json") - var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, mode, true, dstate.OpenDmsArgs{})) - if mode != dstate.WritingFile { - db.InitializeOperationBuffer(ctx, "dep-1", 1, 2) - } - dependsOn := []deployplan.DependsOnEntry{{Node: "resources.pipelines.upstream", Label: "upstream"}} - require.NoError(t, db.SaveState(ctx, "resources.jobs.kept", "keep-id", map[string]string{"name": "kept"}, dependsOn)) - require.NoError(t, db.SaveState(ctx, "resources.jobs.deleted", "delete-id", map[string]string{"name": "deleted"}, nil)) - require.NoError(t, db.DeleteState(ctx, "resources.jobs.deleted", false)) - exported, err := db.Finalize(ctx) - require.NoError(t, err) - assert.Len(t, exported, 1) - assert.Equal(t, "keep-id", exported["resources.jobs.kept"].ID) - assert.NoFileExists(t, path+".wal") - completed, err := db.CompleteVersion(ctx, true) - require.NoError(t, err) - assert.Equal(t, mode != dstate.WritingFile, completed) - completed, err = db.CompleteVersion(ctx, true) - require.NoError(t, err) - assert.False(t, completed) - - var reopened dstate.DeploymentState - require.NoError(t, reopened.Open(t.Context(), path, dstate.ReadingFile, false, dstate.OpenDmsArgs{})) - assert.Equal(t, 1, reopened.GetSerial()) - assert.Len(t, reopened.Data.State, 1) - assert.Equal(t, dependsOn, reopened.Data.State["resources.jobs.kept"].DependsOn) - _, enforced := reopened.Data.Features[dstate.FeatureDeploymentHistoryBestEffort] - assert.Equal(t, mode == dstate.WritingFileWithHistoryEnforced, enforced) - _, err = reopened.Finalize(t.Context()) - require.NoError(t, err) - - if mode == dstate.WritingFile { - assert.Empty(t, updates) - } else { - require.Len(t, updates, 2) - assert.JSONEq(t, `{"state":{"name":"kept"},"depends_on":[{"node":"resources.pipelines.upstream","label":"upstream"}]}`, updates["jobs.kept"]["state"].(string)) - assert.NotContains(t, updates["jobs.deleted"], "state") - } - }) - } -} - -func TestHistoryErrorsFailOpen(t *testing.T) { - for _, mode := range []dstate.Mode{dstate.WritingDMS, dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { - t.Run(fmt.Sprint(mode), func(t *testing.T) { - goodUploads := 0 - ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { - if r.Method == http.MethodPatch && filepath.Base(r.URL.Path) == "jobs.good" { - goodUploads++ - _, _ = w.Write([]byte(`{"sequence_id":1}`)) - return - } - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusBadRequest) - _, _ = w.Write([]byte(`{"error_code":"INVALID_PARAMETER_VALUE","message":"history unavailable"}`)) - }) - path := filepath.Join(t.TempDir(), "resources.json") - var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) - db.InitializeOperationBuffer(ctx, "dep-1", 1, 2) - require.NoError(t, db.SaveState(ctx, "resources.jobs.job", "id-1", map[string]string{"name": "job"}, nil)) - require.NoError(t, db.SaveState(ctx, "resources.jobs.good", "id-2", struct{}{}, nil)) - _, err := db.Finalize(ctx) - assert.Equal(t, 1, goodUploads, "a bad upload must not stop the next resource") - if mode == dstate.WritingDMS { - require.ErrorContains(t, err, "history unavailable") - assert.Error(t, db.RecordingError()) - } else { - require.NoError(t, err) - assert.NoError(t, db.RecordingError()) - var reopened dstate.DeploymentState - require.NoError(t, reopened.Open(t.Context(), path, dstate.ReadingFile, false, dstate.OpenDmsArgs{})) - assert.Equal(t, "id-1", reopened.GetResourceID("resources.jobs.job")) - _, err = reopened.Finalize(t.Context()) - require.NoError(t, err) - } - _, err = db.CompleteVersion(ctx, true) - if mode == dstate.WritingDMS { - assert.ErrorContains(t, err, "history unavailable") - } else { - assert.NoError(t, err) - } - completed, err := db.CompleteVersion(ctx, true) - require.NoError(t, err) - assert.False(t, completed) - }) - } -} - -func TestEnforcedHistorySurvivesWALRecovery(t *testing.T) { - ctx := t.Context() - path := filepath.Join(t.TempDir(), "resources.json") - var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, dstate.WritingFileWithHistoryEnforced, false, dstate.OpenDmsArgs{})) - wal, err := os.ReadFile(path + ".wal") - require.NoError(t, err) - _, err = db.Finalize(ctx) - require.NoError(t, err) - - // Recover an interrupted first deployment: its WAL contains only the feature marker. - recoveredPath := filepath.Join(t.TempDir(), "resources.json") - require.NoError(t, os.WriteFile(recoveredPath+".wal", wal, 0o600)) - var recovered dstate.DeploymentState - require.NoError(t, recovered.Open(ctx, recoveredPath, dstate.ReadingFile, true, dstate.OpenDmsArgs{})) - assert.Contains(t, recovered.StateFeatures(), dstate.FeatureDeploymentHistoryBestEffort) - assert.Equal(t, 1, recovered.GetSerial()) - assert.FileExists(t, recoveredPath) - assert.NoFileExists(t, recoveredPath+".wal") - _, err = recovered.Finalize(ctx) - require.NoError(t, err) -} - -func TestHistoryQueueOnlyTimesOutAtDrain(t *testing.T) { - for _, mode := range []dstate.Mode{dstate.WritingFileWithHistory, dstate.WritingFileWithHistoryEnforced} { - t.Run(fmt.Sprint(mode), func(t *testing.T) { - synctest.Test(t, func(t *testing.T) { - started := make(chan struct{}) - finished := make(chan struct{}) - ctx := modeContext(t, func(w http.ResponseWriter, r *http.Request) { - _, _ = io.Copy(io.Discard, r.Body) - close(started) - <-r.Context().Done() - close(finished) - }) - path := filepath.Join(t.TempDir(), "resources.json") - var db dstate.DeploymentState - require.NoError(t, db.Open(ctx, path, mode, false, dstate.OpenDmsArgs{})) - const operationCount = 32 - db.InitializeOperationBuffer(ctx, "dep-1", 1, operationCount) - start := time.Now() - for i := range operationCount { - key := fmt.Sprintf("resources.jobs.job_%d", i) - require.NoError(t, db.SaveState(ctx, key, strconv.Itoa(i), struct{}{}, nil)) - } - assert.Zero(t, time.Since(start)) - <-started - time.Sleep(time.Minute + time.Second) - select { - case <-finished: - t.Fatal("recording stopped before drain") - default: - } - drainStart := time.Now() - exported, err := db.Finalize(ctx) - require.NoError(t, err) - assert.Equal(t, 160*time.Second, time.Since(drainStart)) - assert.Len(t, exported, operationCount) - assert.FileExists(t, path) - assert.NoFileExists(t, path+".wal") - <-finished - assert.Equal(t, 160*time.Second, time.Since(drainStart)) - }) - }) - } -} diff --git a/bundle/direct/dstate/state_test.go b/bundle/direct/dstate/state_test.go index ce41c2dc7d1..cd2e1bbc1cc 100644 --- a/bundle/direct/dstate/state_test.go +++ b/bundle/direct/dstate/state_test.go @@ -421,3 +421,36 @@ func TestOpenFailureLeavesStateClosed(t *testing.T) { assert.Equal(t, "test-lineage", db.Data.Lineage) mustFinalize(t, &db) } + +func TestGetSerialWithHistoryUsesLatestCounter(t *testing.T) { + for _, mode := range []Mode{WritingFileWithHistory, WritingFileWithHistoryEnforced} { + db := DeploymentState{mode: mode, Data: NewDatabase("", 4), VersionID: 7} + assert.Equal(t, 7, db.GetSerial(), "mode=%d: service version is ahead", mode) + db.Data.Serial = 9 + assert.Equal(t, 9, db.GetSerial(), "mode=%d: file serial is ahead", mode) + } +} + +func TestHeaderOnlyWALPreservesHistoryFeature(t *testing.T) { + path := filepath.Join(t.TempDir(), "state.json") + header := Header{ + Lineage: "lineage", + Serial: 1, + StateVersion: currentStateVersion, + Features: map[string]struct{}{FeatureDeploymentHistoryBestEffort: {}}, + } + wal, err := json.Marshal(header) + require.NoError(t, err) + require.NoError(t, os.WriteFile(path+walSuffix, append(wal, '\n'), 0o600)) + + var db DeploymentState + require.NoError(t, db.Open(t.Context(), path, ReadingFile, WithRecovery(true), OpenDmsArgs{})) + data, err := os.ReadFile(path) + require.NoError(t, err) + var persisted Database + require.NoError(t, json.Unmarshal(data, &persisted)) + assert.Contains(t, persisted.Features, FeatureDeploymentHistoryBestEffort) + assert.Equal(t, 1, persisted.Serial) + assert.NoFileExists(t, path+walSuffix) + mustFinalize(t, &db) +} From dae6556d286ac42da978409cf00ce953920bd959 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 12:27:43 +0000 Subject: [PATCH 47/52] Remove the header-only history feature recovery test --- bundle/direct/dstate/state_test.go | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/bundle/direct/dstate/state_test.go b/bundle/direct/dstate/state_test.go index cd2e1bbc1cc..a44afae8fab 100644 --- a/bundle/direct/dstate/state_test.go +++ b/bundle/direct/dstate/state_test.go @@ -430,27 +430,3 @@ func TestGetSerialWithHistoryUsesLatestCounter(t *testing.T) { assert.Equal(t, 9, db.GetSerial(), "mode=%d: file serial is ahead", mode) } } - -func TestHeaderOnlyWALPreservesHistoryFeature(t *testing.T) { - path := filepath.Join(t.TempDir(), "state.json") - header := Header{ - Lineage: "lineage", - Serial: 1, - StateVersion: currentStateVersion, - Features: map[string]struct{}{FeatureDeploymentHistoryBestEffort: {}}, - } - wal, err := json.Marshal(header) - require.NoError(t, err) - require.NoError(t, os.WriteFile(path+walSuffix, append(wal, '\n'), 0o600)) - - var db DeploymentState - require.NoError(t, db.Open(t.Context(), path, ReadingFile, WithRecovery(true), OpenDmsArgs{})) - data, err := os.ReadFile(path) - require.NoError(t, err) - var persisted Database - require.NoError(t, json.Unmarshal(data, &persisted)) - assert.Contains(t, persisted.Features, FeatureDeploymentHistoryBestEffort) - assert.Equal(t, 1, persisted.Serial) - assert.NoFileExists(t, path+walSuffix) - mustFinalize(t, &db) -} From c60e9d93e10de17ccef0d1b6a884dd038354c718 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 12:36:24 +0000 Subject: [PATCH 48/52] Size history channels and drain budgets from staged operations --- bundle/direct/dstate/state.go | 18 ++----- libs/dms/operation_buffer.go | 88 +++++++++++-------------------- libs/dms/operation_buffer_test.go | 48 ++++++++++------- 3 files changed, 66 insertions(+), 88 deletions(-) diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index dc99ddc2ade..411eeb2af3d 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -14,7 +14,6 @@ import ( "slices" "strings" "sync" - "time" "uuid" "github.com/databricks/cli/bundle/config" @@ -38,8 +37,6 @@ const ( walSuffix = ".wal" ) -const historyDrainTimeoutPerOperation = 5 * time.Second - // FeatureDeploymentHistory marks a state whose resources are also recorded with the // deployment metadata service. Both stores are kept in step, so the marker is what tells a // reader the two already agree. A CLI that does not know the name refuses the state rather @@ -107,8 +104,6 @@ type DeploymentState struct { // history, in which case InitializeOperationBuffer installs it once the version exists. operationBuffer *dms.OperationBuffer - historyDrainTimeout time.Duration - // versionCompleted makes CompleteVersion a no-op after the first call, so a deferred safety-net // completion after an explicit one does nothing. versionCompleted bool @@ -186,11 +181,10 @@ func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploy db.mu.Lock() defer db.mu.Unlock() options := dms.BufferOptions{ - BestEffort: db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced, + BestEffort: db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced, + OperationCount: operationCount, } db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options) - // Use the full plan count; the queue length changes as uploads finish. - db.historyDrainTimeout = time.Duration(operationCount) * historyDrainTimeoutPerOperation db.DeploymentID = deploymentID db.VersionID = versionID } @@ -563,7 +557,6 @@ func (db *DeploymentState) reset() { db.Path = "" db.Data = Database{} db.stateIDs = nil - db.historyDrainTimeout = 0 // Telemetry and destroy still need the mode and IDs after version completion. // Clear the recording buffer once CompleteVersion has copied it. if db.versionCompleted { @@ -888,11 +881,8 @@ func (db *DeploymentState) Finalize(ctx context.Context) (resourcestate.Exported case WritingFileWithHistory, WritingFileWithHistoryEnforced: if db.operationBuffer != nil { // File state is already persisted, so history errors only warn. - // Bound the wait so a stalled service cannot hold up deployment. - drainCtx, cancel := context.WithTimeout(ctx, db.historyDrainTimeout) - drainErr := db.operationBuffer.Drain(drainCtx) - cancel() - if drainErr != nil { + // The buffer bounds the wait using the staged operation count. + if drainErr := db.operationBuffer.Drain(ctx); drainErr != nil { log.Warnf(ctx, "Failed to record deployment history: %s", drainErr) } } diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 8e8973ecf7d..615f7802ca7 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -15,9 +15,14 @@ import ( // plan reads. const bufferedOperations = 10 +const historyDrainTimeoutPerOperation = 5 * time.Second + // BufferOptions selects optional file-history recording behavior. type BufferOptions struct { BestEffort bool + // OperationCount counts staged resource keys. All recorded keys must be staged. + // File history uses it for queue capacity and the drain timeout. + OperationCount int } // OperationBuffer records each state write with DMS for one deployment version, off the apply @@ -34,8 +39,9 @@ type OperationBuffer struct { stopQueue func() // cancel stops in-flight history uploads when the drain deadline expires. - cancel context.CancelFunc - bestEffort bool + cancel context.CancelFunc + bestEffort bool + drainTimeout time.Duration // sequenceIDs holds the token the last update for a resource returned. A missing key reads as // 0, which is what staging leaves. Unguarded: run is the only goroutine that writes, one update at a time. @@ -49,13 +55,6 @@ type OperationBuffer struct { // mu guards the fields below. mu sync.Mutex - // File-history recording uses an unbounded queue so DMS applies no backpressure - // to the deployment engine. - historyQueue []string - - // draining lets the worker exit once the history queue is empty. - draining bool - // pending holds the newest update per resource key, absent once the writer takes it. pending map[string]OperationUpdate @@ -73,13 +72,21 @@ type OperationBuffer struct { func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum int, options BufferOptions) *OperationBuffer { // History can stop independently without canceling the deployment. ctx, cancel := context.WithCancel(ctx) + capacity := bufferedOperations + if options.BestEffort { + // Updates coalesce, so each staged key occupies at most one queue slot. + // The worker frees that slot before the key can be queued again. + // One slot per staged key therefore avoids deployment backpressure. + capacity = options.OperationCount + } b := &OperationBuffer{ deploymentID: deploymentID, versionNum: versionNum, - queue: make(chan string, bufferedOperations), + queue: make(chan string, capacity), done: make(chan struct{}), cancel: cancel, bestEffort: options.BestEffort, + drainTimeout: time.Duration(options.OperationCount) * historyDrainTimeoutPerOperation, pending: make(map[string]OperationUpdate), latestState: make(map[string]OperationUpdate), sequenceIDs: make(map[string]int64), @@ -126,13 +133,6 @@ func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, kee update = waiting.Merge(update) } b.pending[resourceKey] = update - if b.bestEffort { - if !queued { - b.historyQueue = append(b.historyQueue, resourceKey) - } - b.mu.Unlock() - return - } b.mu.Unlock() // when already queued: the writer reads the map when it gets to the key, @@ -160,44 +160,20 @@ func (b *OperationBuffer) run(ctx context.Context) { for { var resourceKey string + var ok bool if b.bestEffort { - // Calling b.cancel() makes ctx.Err() return an error. This is the signal to stop. - // We expect this after a drain timeout. - // Cancellation or expiry of the deployment's context also stops the worker here. - if err := ctx.Err(); err != nil { - b.setErr(err) + select { + case resourceKey, ok = <-b.queue: + case <-ctx.Done(): + // Drain calls cancel on timeout. Parent cancellation can also stop history. + b.setErr(ctx.Err()) return } - b.mu.Lock() - if len(b.historyQueue) == 0 { - draining := b.draining - b.mu.Unlock() - if draining { - return - } - // Start with polling to keep the design simple. - // We can optimize this later by waking the worker when work is queued. - // Cancellation ends this wait immediately. - select { - case <-time.After(100 * time.Millisecond): - case <-ctx.Done(): - } - continue - } - resourceKey = b.historyQueue[0] - b.historyQueue[0] = "" - b.historyQueue = b.historyQueue[1:] - if len(b.historyQueue) == 0 { - // Release the backing array so the garbage collector can reclaim it. - b.historyQueue = nil - } - b.mu.Unlock() } else { - var ok bool resourceKey, ok = <-b.queue - if !ok { - return - } + } + if !ok { + return } update, ok := b.take(resourceKey) if !ok { @@ -232,21 +208,21 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio // Drain waits for buffered uploads before version completion. // - DMS state: close the channel and wait for queued writes. -// - File history: finish the queue, or cancel uploads when ctx expires. +// - File history: wait up to five seconds per staged operation, then cancel uploads. // // It returns the first recording error and is safe to call more than once. func (b *OperationBuffer) Drain(ctx context.Context) error { + b.stopQueue() if !b.bestEffort { // DMS is authoritative: wait for all queued writes and return any recording error. - b.stopQueue() <-b.done return b.Err() } - // File state is already saved. Try the remaining history writes until ctx expires. - b.mu.Lock() - b.draining = true - b.mu.Unlock() + // File state is already saved. Budget five seconds per staged operation, + // including operations already uploaded; caller cancellation can shorten the wait. + ctx, cancel := context.WithTimeout(ctx, b.drainTimeout) + defer cancel() select { case <-b.done: // Report the first upload error even if later uploads succeeded. diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index 14d17aa303b..208856a2c4d 100644 --- a/libs/dms/operation_buffer_test.go +++ b/libs/dms/operation_buffer_test.go @@ -6,8 +6,10 @@ import ( "errors" "fmt" "strings" + "sync" "testing" "testing/synctest" + "time" "github.com/databricks/databricks-sdk-go/service/bundledeployments" "github.com/stretchr/testify/assert" @@ -112,33 +114,43 @@ func TestRecordFailureUsesTheGivenStateWhenNoWriteRecorded(t *testing.T) { assert.Equal(t, bundledeployments.OperationStatusOperationStatusFailed, update.Status) } -func TestBufferQueuesWithoutBackpressure(t *testing.T) { +func TestHistoryQueueDoesNotBlockWhenWriterStops(t *testing.T) { synctest.Test(t, func(t *testing.T) { - b := queued() - b.bestEffort = true const operationCount = 1000 - for i := range operationCount { - key := fmt.Sprintf("resources.jobs.job_%d", i) - b.RecordOperation(t.Context(), key, false, "id-1", json.RawMessage(`{"state":{}}`)) + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + b := StartOperationBuffer(ctx, "dep-1", 1, BufferOptions{BestEffort: true, OperationCount: operationCount}) + synctest.Wait() + cancel() + <-b.done + assert.ErrorIs(t, b.Err(), context.Canceled) + + // Even with no consumer, every staged key fits. Repeated writes coalesce. + for range 2 { + for i := range operationCount { + key := fmt.Sprintf("resources.jobs.job_%d", i) + b.RecordOperation(ctx, key, false, "id-1", json.RawMessage(`{"state":{}}`)) + } } - assert.NoError(t, b.Err()) - assert.Len(t, b.historyQueue, operationCount) + assert.Len(t, b.queue, operationCount) assert.Len(t, b.pending, operationCount) - for i, key := range b.historyQueue { - assert.Equal(t, fmt.Sprintf("resources.jobs.job_%d", i), key) - } }) } -func TestBufferStopsIdleWriterOnCancellation(t *testing.T) { +func TestHistoryDrainCancelsStalledWriter(t *testing.T) { synctest.Test(t, func(t *testing.T) { ctx, cancel := context.WithCancel(t.Context()) defer cancel() - b := StartOperationBuffer(ctx, "dep-1", 1, BufferOptions{BestEffort: true}) - synctest.Wait() - cancel() - <-b.done - assert.ErrorIs(t, b.Err(), context.Canceled) + b := queued() + b.bestEffort = true + b.cancel = cancel + b.drainTimeout = time.Second + b.done = make(chan struct{}) // Simulate a writer that cannot finish its upload. + b.stopQueue = sync.OnceFunc(func() { close(b.queue) }) + start := time.Now() + assert.ErrorIs(t, b.Drain(t.Context()), context.DeadlineExceeded) + assert.Equal(t, time.Second, time.Since(start)) + assert.ErrorIs(t, ctx.Err(), context.Canceled) }) } @@ -149,6 +161,6 @@ func TestBestEffortBufferKeepsRecordingAfterError(t *testing.T) { b.RecordOperation(t.Context(), "resources.jobs.bar", false, "id-2", json.RawMessage(`{}`)) b.RecordFailure("resources.jobs.bar", "id-2", nil, errors.New("apply failed")) assert.ErrorContains(t, b.Err(), "exceeds") - assert.Len(t, b.historyQueue, 1) + assert.Len(t, b.queue, 1) assert.Contains(t, b.pending, "resources.jobs.bar") } From ab75f6e611dfe7887d1f468b7527fb59ae101adc Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 12:40:12 +0000 Subject: [PATCH 49/52] Remove redundant locking detail from record comment --- libs/dms/operation_buffer.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 615f7802ca7..d043912c422 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -118,7 +118,7 @@ func (b *OperationBuffer) RecordFailure(resourceKey, resourceID string, state js // record queues update for resourceKey. keepExistingUpdateIfAny (a failure) merges onto the // resource's last recorded state if one exists; otherwise (a state write) update is saved as -// that latest state. latestState and pending are updated under the one lock. +// that latest state. func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, keepExistingUpdateIfAny bool, cause error) { b.mu.Lock() if keepExistingUpdateIfAny { From 3d75aa8b386255576248eb4ba0e9b3ce095fbb06 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 12:41:08 +0000 Subject: [PATCH 50/52] Name deployment state modes explicitly in Drain documentation --- libs/dms/operation_buffer.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index d043912c422..5d25eecc74b 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -207,9 +207,11 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio } // Drain waits for buffered uploads before version completion. -// - DMS state: close the channel and wait for queued writes. -// - File history: wait up to five seconds per staged operation, then cancel uploads. +// - WritingDMS: close the channel and wait for queued writes. +// - WritingFileWithHistory: wait up to five seconds per staged operation. +// - WritingFileWithHistoryEnforced: wait up to five seconds per staged operation. // +// Both history modes cancel outstanding uploads when the timeout or ctx expires. // It returns the first recording error and is safe to call more than once. func (b *OperationBuffer) Drain(ctx context.Context) error { b.stopQueue() From 529f38c38dbf1eb55be8e358df21863c7dba69b1 Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 12:43:42 +0000 Subject: [PATCH 51/52] Explain the stalled writer drain test setup and assertions --- libs/dms/operation_buffer_test.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index 208856a2c4d..d67a11de2a1 100644 --- a/libs/dms/operation_buffer_test.go +++ b/libs/dms/operation_buffer_test.go @@ -138,18 +138,25 @@ func TestHistoryQueueDoesNotBlockWhenWriterStops(t *testing.T) { } func TestHistoryDrainCancelsStalledWriter(t *testing.T) { + // A stalled upload must time out and cancel its worker. + // Fake time lets us test the one-second timeout without a real one-second wait. synctest.Test(t, func(t *testing.T) { + // Track the writer's cancellation separately from Drain's deadline. ctx, cancel := context.WithCancel(t.Context()) defer cancel() b := queued() b.bestEffort = true b.cancel = cancel b.drainTimeout = time.Second - b.done = make(chan struct{}) // Simulate a writer that cannot finish its upload. + // No worker runs or closes done. Drain must eventually time out. + b.done = make(chan struct{}) b.stopQueue = sync.OnceFunc(func() { close(b.queue) }) + + // Drain reports its timeout after exactly one fake second. start := time.Now() assert.ErrorIs(t, b.Drain(t.Context()), context.DeadlineExceeded) assert.Equal(t, time.Second, time.Since(start)) + // Drain must also call cancel(), stopping any upload the writer would be doing. assert.ErrorIs(t, ctx.Err(), context.Canceled) }) } From 1723931ec92839cd9463055b75f061e44132ed1e Mon Sep 17 00:00:00 2001 From: Shreyas Goenka Date: Fri, 23 Oct 2026 12:46:03 +0000 Subject: [PATCH 52/52] Explain fake clock behavior in the history drain test --- libs/dms/operation_buffer_test.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index d67a11de2a1..0e019284360 100644 --- a/libs/dms/operation_buffer_test.go +++ b/libs/dms/operation_buffer_test.go @@ -139,7 +139,9 @@ func TestHistoryQueueDoesNotBlockWhenWriterStops(t *testing.T) { func TestHistoryDrainCancelsStalledWriter(t *testing.T) { // A stalled upload must time out and cancel its worker. - // Fake time lets us test the one-second timeout without a real one-second wait. + // synctest runs this function with a fake clock. When the test blocks waiting + // for Drain, it advances time to the next timer: the one-second drain deadline. + // No real second passes, and time.Since(start) reports exactly one second. synctest.Test(t, func(t *testing.T) { // Track the writer's cancellation separately from Drain's deadline. ctx, cancel := context.WithCancel(t.Context())