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..c7d6440573c --- /dev/null +++ b/bundle/direct/dstate/mode.go @@ -0,0 +1,36 @@ +package dstate + +// Mode selects the source of resource state and the persistence required for writes. +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 + + // WritingDMS reads and writes resource state through 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 deployment_history_best_effort to reject clients that cannot honor history recording. + WritingFileWithHistoryEnforced +) + +func (m Mode) isWritable() bool { + switch m { + case WritingFile, WritingDMS, WritingFileWithHistory, WritingFileWithHistoryEnforced: + return true + default: + return false + } +} diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 4d2a82546f9..411eeb2af3d 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -47,6 +47,10 @@ const ( // deploying such a state without recording is refused. const FeatureDeploymentHistory = "deployment_history" +// 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. const featuresDocURL = "https://docs.databricks.com/aws/en/dev-tools/bundles/state-features#state-features" @@ -55,7 +59,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: {}, + FeatureDeploymentHistory: {}, + FeatureDeploymentHistoryBestEffort: {}, } // assertNoUnsupportedFeatures refuses a state that depends on a feature this CLI does not recognize. @@ -88,9 +93,9 @@ 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 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). stateIDs map[string]string @@ -103,11 +108,6 @@ 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 - // 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 @@ -177,12 +177,14 @@ 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.mu.Lock() defer db.mu.Unlock() - db.operationBuffer = buf + options := dms.BufferOptions{ + BestEffort: db.mode == WritingFileWithHistory || db.mode == WritingFileWithHistoryEnforced, + OperationCount: operationCount, + } + db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options) db.DeploymentID = deploymentID db.VersionID = versionID } @@ -195,9 +197,14 @@ 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 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() if buf == nil { return nil @@ -211,7 +218,7 @@ 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() { + if !db.RecordsHistory() { return false, nil } @@ -223,6 +230,8 @@ 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 + db.reset() db.mu.Unlock() // A recording failure fails the version even when the caller counted the deploy a success: the @@ -242,6 +251,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 +266,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,10 +320,14 @@ 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() { + switch db.mode { + case WritingFile: + 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 if buf := db.operationBuffer; buf != nil { @@ -321,13 +338,25 @@ func (db *DeploymentState) SaveState(ctx context.Context, key, newID string, sta buf.RecordOperation(ctx, key, false, newID, recorded) } return nil - } - err = appendJSONLine(db.walFile, WALEntry{Key: key, Value: &entry}) - if err == nil { + 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) + } + return nil + + default: + return fmt.Errorf("cannot save state in mode %d", db.mode) } - return err } // DeleteState drops the resource's state entry: the resource is gone. inProgress records the @@ -344,17 +373,34 @@ 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() { + switch db.mode { + case WritingFile: + 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 { buf.RecordOperation(ctx, key, inProgress, deletedID, nil) } - } else if err := appendJSONLine(db.walFile, WALEntry{Key: key}); err != nil { - return err - } + delete(db.stateIDs, key) + return nil - delete(db.stateIDs, key) - return nil + 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) + return nil + + default: + return fmt.Errorf("cannot delete state in mode %d", db.mode) + } } func (db *DeploymentState) GetResourceEntry(key string) (ResourceEntry, bool) { @@ -401,21 +447,25 @@ 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 WritingDMS: return db.VersionID + case WritingFileWithHistory, WritingFileWithHistoryEnforced: + return max(db.VersionID, db.Data.Serial) + default: + return db.Data.Serial } - return db.Data.Serial } // isDeploymentMetadataService is IsDeploymentMetadataService for callers already holding db.mu. func (db *DeploymentState) isDeploymentMetadataService() bool { - return db.recordsHistory + return db.mode == WritingDMS } // IsDeploymentMetadataService reports whether this deployment's resource state lives in the @@ -427,6 +477,18 @@ 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() + switch db.mode { + case WritingDMS, WritingFileWithHistory, WritingFileWithHistoryEnforced: + return true + default: + return false + } +} + // 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 +509,27 @@ 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 -) - -// 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 { +// 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 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. +// - 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() @@ -475,7 +537,7 @@ 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 @@ -495,14 +557,19 @@ func (db *DeploymentState) reset() { db.Path = "" db.Data = Database{} db.stateIDs = nil - db.openedForWrite = false + // 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 + } } -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 { 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 = mode + 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,8 +595,6 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, withRe db.stateIDs[key] = entry.ID } - recording := bool(withDeploymentHistory) - walPath := db.Path + walSuffix _, err = os.Stat(walPath) switch { @@ -539,7 +604,7 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, withRe 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. @@ -558,26 +623,24 @@ func (db *DeploymentState) unlockedOpen(ctx context.Context, path string, withRe // 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 } - 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,46 +675,14 @@ 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) - } + switch mode { + case ReadingFile: + return nil + case WritingDMS: + return nil + case WritingFile, WritingFileWithHistory, WritingFileWithHistoryEnforced: + return db.setupWal() } - return nil } @@ -666,6 +697,7 @@ func (db *DeploymentState) OpenWithData(path string, data Database) { } db.Path = path + db.mode = ReadingFile db.Data = data db.stateIDs = make(map[string]string) for key, entry := range data.State { @@ -714,6 +746,7 @@ func (db *DeploymentState) mergeWalIntoState(ctx context.Context) (bool, error) var ( newSerial int newCLIVersion string + newFeatures map[string]struct{} ) for scanner.Scan() { @@ -738,6 +771,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) } + newFeatures = header.Features newSerial = header.Serial newCLIVersion = header.CLIVersion } else { @@ -775,6 +809,11 @@ func (db *DeploymentState) mergeWalIntoState(ctx context.Context) (bool, error) } 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 entries, because the caller // (replayWAL) persists the new state file only in that case. A header-only @@ -790,6 +829,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 +855,12 @@ 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 { + // 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 @@ -824,9 +869,22 @@ 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(); 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) + } + } + case WritingFileWithHistory, WritingFileWithHistoryEnforced: + if db.operationBuffer != nil { + // File state is already persisted, so history errors only warn. + // 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) + } } } @@ -837,8 +895,8 @@ 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. +// 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() @@ -846,17 +904,27 @@ func (db *DeploymentState) UpgradeToWrite() error { if db.Path == "" { return errors.New("internal error: DeploymentState must be opened first") } - if db.openedForWrite { + switch db.mode { + case WritingDMS, WritingFileWithHistory, WritingFileWithHistoryEnforced: + return nil + case WritingFile: return errors.New("internal error: DeploymentState is already open for write") + case ReadingFile: + db.mode = WritingFile + if err := db.setupWal(); err != nil { + db.mode = ReadingFile + return err + } + return nil } - db.openedForWrite = true + return fmt.Errorf("invalid deployment state mode: %d", db.mode) +} - // As in Open: a recorded deployment writes no WAL, so stamp the header and skip the file. - if db.isDeploymentMetadataService() { - db.Data.Lineage = db.GetOrInitLineage() - db.Data.StateVersion = currentStateVersion - db.Data.CLIVersion = build.GetInfo().Version - return nil +// 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 } walPath := db.Path + walSuffix @@ -868,14 +936,30 @@ func (db *DeploymentState) UpgradeToWrite() error { 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, + Features: maps.Clone(db.Data.Features), } - return appendJSONLine(db.walFile, walHead) + 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 + // 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) + } + return nil } // Persist writes the current in-memory state to its file immediately, without a WAL cycle. @@ -903,14 +987,16 @@ func (db *DeploymentState) AssertOpenedForReadOrWrite() { func (db *DeploymentState) AssertOpenedForRead() { db.AssertOpenedForReadOrWrite() - if db.openedForWrite { + // 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") } } 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..a44afae8fab 100644 --- a/bundle/direct/dstate/state_test.go +++ b/bundle/direct/dstate/state_test.go @@ -25,14 +25,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 +42,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 +63,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 = WritingDMS } db.Data.State["resources.jobs.foo"] = ResourceEntry{ID: "123", State: json.RawMessage(`{"name":"foo"}`)} @@ -163,6 +166,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{}, @@ -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, WritingDMS, WithRecovery(false), OpenDmsArgs{ DeploymentID: tt.deploymentID, LastVersionID: tt.lastVersionID, }) if tt.wantErr != "" { @@ -189,6 +193,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) }) } @@ -243,10 +248,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 +269,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 +296,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 +308,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 +329,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 +358,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 +380,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 @@ -392,7 +397,7 @@ func TestGetOrInitLineageReadableBeforeWriteAndPersisted(t *testing.T) { // 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 +411,22 @@ 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 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) + } +} 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..4fefd8504ac 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) } diff --git a/cmd/bundle/utils/process.go b/cmd/bundle/utils/process.go index 21c3008f719..8fa68f32f5e 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.WritingDMS, 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.WritingDMS, 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..5d25eecc74b 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" @@ -14,6 +15,16 @@ 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 // 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. @@ -27,6 +38,11 @@ type OperationBuffer struct { done chan struct{} stopQueue func() + // cancel stops in-flight history uploads when the drain deadline expires. + 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. // @@ -53,12 +69,24 @@ 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 { + // 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), @@ -90,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. Panics after Drain. +// that latest state. func (b *OperationBuffer) record(resourceKey string, update OperationUpdate, keepExistingUpdateIfAny bool, cause error) { b.mu.Lock() if keepExistingUpdateIfAny { @@ -128,8 +156,25 @@ 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 { + for { + var resourceKey string + var ok bool + if b.bestEffort { + 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 + } + } else { + resourceKey, ok = <-b.queue + } + if !ok { + return + } update, ok := b.take(resourceKey) if !ok { // Unreachable: a key is queued only when nothing was waiting for it. Guard so @@ -161,15 +206,38 @@ 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 uploads before version completion. +// - 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() - <-b.done - return b.Err() + if !b.bestEffort { + // DMS is authoritative: wait for all queued writes and return any recording error. + <-b.done + return b.Err() + } + + // 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. + return b.Err() + case <-ctx.Done(): + // 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() diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index 49af18c3255..0e019284360 100644 --- a/libs/dms/operation_buffer_test.go +++ b/libs/dms/operation_buffer_test.go @@ -1,11 +1,15 @@ package dms import ( + "context" "encoding/json" "errors" + "fmt" "strings" "sync" "testing" + "testing/synctest" + "time" "github.com/databricks/databricks-sdk-go/service/bundledeployments" "github.com/stretchr/testify/assert" @@ -69,15 +73,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()) - require.NoError(t, s.Drain()) + // 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())) } func TestRecordFailureKeepsTheStateTheLastWriteRecorded(t *testing.T) { @@ -114,3 +113,63 @@ func TestRecordFailureUsesTheGivenStateWhenNoWriteRecorded(t *testing.T) { assert.JSONEq(t, `{"state":{"name":"v1"}}`, string(update.State)) assert.Equal(t, bundledeployments.OperationStatusOperationStatusFailed, update.Status) } + +func TestHistoryQueueDoesNotBlockWhenWriterStops(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + const operationCount = 1000 + 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.Len(t, b.queue, operationCount) + assert.Len(t, b.pending, operationCount) + }) +} + +func TestHistoryDrainCancelsStalledWriter(t *testing.T) { + // A stalled upload must time out and cancel its worker. + // 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()) + defer cancel() + b := queued() + b.bestEffort = true + b.cancel = cancel + b.drainTimeout = time.Second + // 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) + }) +} + +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.Len(t, b.queue, 1) + assert.Contains(t, b.pending, "resources.jobs.bar") +}