From dcf4b3c68c9cc7500502d9427a4897c7f376d4ed Mon Sep 17 00:00:00 2001 From: Sankalp-Mittal Date: Wed, 4 Nov 2026 12:27:42 +0000 Subject: [PATCH 1/4] Send heartbeats for open DMS deployment versions DMS marks a version in progress at CreateVersion. If the CLI dies before CompleteVersion, nothing refreshes the version and the UI shows it as deploying forever. The service will expose a per-version last update time, refreshed by UpdateOperation and Heartbeat, and the UI will show a version as interrupted once that time is over two minutes old. The CLI now keeps that time fresh while a version is open. A heartbeat is sent only after 30s with no successful UpdateOperation, so a fast deploy sends none. Heartbeats are best effort: a failed one is retried quickly, then reported as a warning, then tried again at the next interval. It never fails or stops the deploy. Co-authored-by: Isaac --- bundle/direct/dstate/state.go | 11 +++- libs/dms/heartbeat.go | 121 ++++++++++++++++++++++++++++++++++ libs/dms/heartbeat_test.go | 115 ++++++++++++++++++++++++++++++++ libs/dms/operation_buffer.go | 7 +- 4 files changed, 252 insertions(+), 2 deletions(-) create mode 100644 libs/dms/heartbeat.go create mode 100644 libs/dms/heartbeat_test.go diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index ac8f8226c4b..8a74b072b2d 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -98,6 +98,9 @@ type DeploymentState struct { // history, in which case InitializeOperationBuffer installs it once the version exists. operationBuffer *dms.OperationBuffer + // heartbeat keeps the open version's last update time fresh. Set together with operationBuffer. + heartbeat *dms.Heartbeat + // versionCompleted makes CompleteVersion a no-op after the first call, so a deferred safety-net // completion after an explicit one does nothing. versionCompleted bool @@ -177,10 +180,12 @@ 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 int) { - buf := dms.StartOperationBuffer(ctx, deploymentID, versionID) + heartbeat := dms.StartHeartbeat(ctx, deploymentID, versionID) + buf := dms.StartOperationBuffer(ctx, deploymentID, versionID, heartbeat) db.mu.Lock() defer db.mu.Unlock() + db.heartbeat = heartbeat db.operationBuffer = buf db.DeploymentID = deploymentID db.VersionID = versionID @@ -222,8 +227,12 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b } db.versionCompleted = true deploymentID, versionID := db.DeploymentID, db.VersionID + heartbeat := db.heartbeat db.mu.Unlock() + // Stopped first: a heartbeat racing the completion would be rejected and warn for nothing. + heartbeat.Stop() + // 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/heartbeat.go b/libs/dms/heartbeat.go new file mode 100644 index 00000000000..73a57b28a95 --- /dev/null +++ b/libs/dms/heartbeat.go @@ -0,0 +1,121 @@ +package dms + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/logdiag" + "github.com/databricks/databricks-sdk-go/service/bundledeployments" +) + +// heartbeatInterval is how long a version may go without an update before a heartbeat is sent. +// The UI shows a version as interrupted once its last update is over 2 minutes old. +const heartbeatInterval = 30 * time.Second + +// heartbeatRetryDelays are the quick retries after a failed heartbeat, before waiting a full interval. +var heartbeatRetryDelays = []time.Duration{time.Second, 2 * time.Second, 4 * time.Second} + +// Heartbeat keeps an open deployment version's last update time fresh, which is how the UI tells +// a running deploy from an interrupted one. CreateVersion and every successful write update it +// too, so a heartbeat is sent only after an interval without one. Best effort: a failure can only +// make the deploy show as interrupted, so it never fails the deploy. +type Heartbeat struct { + send func(ctx context.Context) error + + renewed chan struct{} + stop chan struct{} + done chan struct{} + stopOnce func() +} + +// StartHeartbeat starts heartbeating the version the caller just created. CreateVersion counts +// as the first update, so the first heartbeat goes out an interval later. +func StartHeartbeat(ctx context.Context, deploymentID string, versionNum int) *Heartbeat { + name := VersionName(deploymentID, versionNum) + send := func(ctx context.Context) error { + _, err := cmdctx.WorkspaceClient(ctx).BundleDeployments.Heartbeat(ctx, bundledeployments.HeartbeatRequest{Name: name}) + return err + } + return startHeartbeat(ctx, send) +} + +func startHeartbeat(ctx context.Context, send func(ctx context.Context) error) *Heartbeat { + h := &Heartbeat{ + send: send, + renewed: make(chan struct{}, 1), + stop: make(chan struct{}), + done: make(chan struct{}), + } + h.stopOnce = sync.OnceFunc(func() { + close(h.stop) + <-h.done + }) + go h.run(ctx) + return h +} + +// renew records that a successful write updated the version, which the service counts as a +// heartbeat, so the next heartbeat is pushed back a full interval. +func (h *Heartbeat) renew() { + select { + case h.renewed <- struct{}{}: + default: + } +} + +func (h *Heartbeat) run(ctx context.Context) { + defer close(h.done) + + next := time.NewTimer(heartbeatInterval) + defer next.Stop() + for { + select { + case <-ctx.Done(): + return + case <-h.stop: + return + case <-h.renewed: + next.Reset(heartbeatInterval) + case <-next.C: + h.beat(ctx) + next.Reset(heartbeatInterval) + } + } +} + +// beat sends a heartbeat, retrying quickly a few times before leaving it to the next interval. +func (h *Heartbeat) beat(ctx context.Context) { + err := h.send(ctx) + for _, delay := range heartbeatRetryDelays { + if err == nil || ctx.Err() != nil { + break + } + log.Warnf(ctx, "Failed to send deployment heartbeat, retrying in %s: %s", delay, err) + select { + case <-time.After(delay): + case <-ctx.Done(): + return + case <-h.stop: + return + } + err = h.send(ctx) + } + + // A cancelled run needs no warning. + if err != nil && ctx.Err() == nil { + logdiag.LogDiag(ctx, diag.Diagnostic{ + Severity: diag.Warning, + Summary: fmt.Sprintf("Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err), + }) + } +} + +// Stop ends the heartbeat and waits for it. Safe to call twice. +func (h *Heartbeat) Stop() { + h.stopOnce() +} diff --git a/libs/dms/heartbeat_test.go b/libs/dms/heartbeat_test.go new file mode 100644 index 00000000000..b7d71c4b2f9 --- /dev/null +++ b/libs/dms/heartbeat_test.go @@ -0,0 +1,115 @@ +package dms + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/databricks/cli/libs/logdiag" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Timing is not reachable from acceptance tests, so the heartbeat is tested here on synctest's +// fake clock, against a send that counts calls instead of the service. + +type fakeSend struct { + calls atomic.Int32 + err error +} + +func (f *fakeSend) send(context.Context) error { + f.calls.Add(1) + return f.err +} + +func heartbeatContext(t *testing.T) context.Context { + ctx := logdiag.InitContext(t.Context()) + logdiag.SetCollect(ctx, true) + return ctx +} + +func TestHeartbeatBeatsWhileIdle(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + f := &fakeSend{} + h := startHeartbeat(heartbeatContext(t), f.send) + defer h.Stop() + + // CreateVersion counts as the first update, so nothing goes out before the interval. + time.Sleep(heartbeatInterval - time.Second) + synctest.Wait() + assert.Equal(t, int32(0), f.calls.Load()) + + time.Sleep(2 * time.Second) + synctest.Wait() + assert.Equal(t, int32(1), f.calls.Load()) + + time.Sleep(heartbeatInterval) + synctest.Wait() + assert.Equal(t, int32(2), f.calls.Load()) + }) +} + +func TestHeartbeatSkippedAfterWrite(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + f := &fakeSend{} + h := startHeartbeat(heartbeatContext(t), f.send) + defer h.Stop() + + time.Sleep(20 * time.Second) + h.renew() + + // The write counts as an update, so the beat due at 30s moves to 50s. + time.Sleep(25 * time.Second) + synctest.Wait() + assert.Equal(t, int32(0), f.calls.Load()) + + time.Sleep(10 * time.Second) + synctest.Wait() + assert.Equal(t, int32(1), f.calls.Load()) + }) +} + +func TestHeartbeatRetriesThenWaitsForInterval(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := heartbeatContext(t) + f := &fakeSend{err: errors.New("connection refused")} + h := startHeartbeat(ctx, f.send) + defer h.Stop() + + // The beat at 30s and its retries 1s, 2s and 4s apart. + time.Sleep(heartbeatInterval + 8*time.Second) + synctest.Wait() + assert.Equal(t, int32(1+len(heartbeatRetryDelays)), f.calls.Load()) + + diags := logdiag.FlushCollected(ctx) + require.Len(t, diags, 1) + assert.Equal(t, "Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in 30s: connection refused", diags[0].Summary) + + // Then nothing more until a full interval has passed. + time.Sleep(heartbeatInterval - 2*time.Second) + synctest.Wait() + assert.Equal(t, int32(1+len(heartbeatRetryDelays)), f.calls.Load()) + + // Checked before its first retry, 1s later. + time.Sleep(1500 * time.Millisecond) + synctest.Wait() + assert.Equal(t, int32(2+len(heartbeatRetryDelays)), f.calls.Load()) + }) +} + +func TestHeartbeatStop(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + f := &fakeSend{} + h := startHeartbeat(heartbeatContext(t), f.send) + + h.Stop() + h.Stop() + time.Sleep(10 * time.Minute) + synctest.Wait() + assert.Equal(t, int32(0), f.calls.Load()) + }) +} diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index efecb3a330a..cd80ab5e40b 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -21,6 +21,9 @@ type OperationBuffer struct { deploymentID string versionNum int + // heartbeat is told of each successful write, which the service counts as a heartbeat. + heartbeat *Heartbeat + // 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 @@ -53,10 +56,11 @@ 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, heartbeat *Heartbeat) *OperationBuffer { b := &OperationBuffer{ deploymentID: deploymentID, versionNum: versionNum, + heartbeat: heartbeat, queue: make(chan string, bufferedOperations), done: make(chan struct{}), pending: make(map[string]OperationUpdate), @@ -158,6 +162,7 @@ func (b *OperationBuffer) write(ctx context.Context, key string, update Operatio // The next write for this resource echoes the sequence id this one earned. b.sequenceIDs[key] = result.SequenceId + b.heartbeat.renew() return nil } From 57340f2b39431cc6d69760a594a2cda05b915751 Mon Sep 17 00:00:00 2001 From: Sankalp-Mittal Date: Wed, 4 Nov 2026 15:24:47 +0000 Subject: [PATCH 2/4] Simplify the heartbeat loop Stop is a context cancel instead of a stop channel, a done channel and a sync.OnceFunc guard. The quick retries after a failed heartbeat are gone: the SDK already retries transient errors inside each call, and the next attempt is at most 30s away against a 2-minute UI threshold. Co-authored-by: Isaac --- libs/dms/heartbeat.go | 103 +++++++++++-------------------------- libs/dms/heartbeat_test.go | 18 ++----- 2 files changed, 35 insertions(+), 86 deletions(-) diff --git a/libs/dms/heartbeat.go b/libs/dms/heartbeat.go index 73a57b28a95..8a2da222ce6 100644 --- a/libs/dms/heartbeat.go +++ b/libs/dms/heartbeat.go @@ -3,12 +3,10 @@ package dms import ( "context" "fmt" - "sync" "time" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/diag" - "github.com/databricks/cli/libs/log" "github.com/databricks/cli/libs/logdiag" "github.com/databricks/databricks-sdk-go/service/bundledeployments" ) @@ -17,105 +15,64 @@ import ( // The UI shows a version as interrupted once its last update is over 2 minutes old. const heartbeatInterval = 30 * time.Second -// heartbeatRetryDelays are the quick retries after a failed heartbeat, before waiting a full interval. -var heartbeatRetryDelays = []time.Duration{time.Second, 2 * time.Second, 4 * time.Second} - // Heartbeat keeps an open deployment version's last update time fresh, which is how the UI tells // a running deploy from an interrupted one. CreateVersion and every successful write update it // too, so a heartbeat is sent only after an interval without one. Best effort: a failure can only -// make the deploy show as interrupted, so it never fails the deploy. +// make the deploy show as interrupted, so it is reported and the next interval tries again. type Heartbeat struct { - send func(ctx context.Context) error - - renewed chan struct{} - stop chan struct{} - done chan struct{} - stopOnce func() + renewed chan struct{} + cancel context.CancelFunc } // StartHeartbeat starts heartbeating the version the caller just created. CreateVersion counts // as the first update, so the first heartbeat goes out an interval later. func StartHeartbeat(ctx context.Context, deploymentID string, versionNum int) *Heartbeat { name := VersionName(deploymentID, versionNum) - send := func(ctx context.Context) error { + return startHeartbeat(ctx, func(ctx context.Context) error { _, err := cmdctx.WorkspaceClient(ctx).BundleDeployments.Heartbeat(ctx, bundledeployments.HeartbeatRequest{Name: name}) return err - } - return startHeartbeat(ctx, send) -} - -func startHeartbeat(ctx context.Context, send func(ctx context.Context) error) *Heartbeat { - h := &Heartbeat{ - send: send, - renewed: make(chan struct{}, 1), - stop: make(chan struct{}), - done: make(chan struct{}), - } - h.stopOnce = sync.OnceFunc(func() { - close(h.stop) - <-h.done }) - go h.run(ctx) - return h } -// renew records that a successful write updated the version, which the service counts as a -// heartbeat, so the next heartbeat is pushed back a full interval. -func (h *Heartbeat) renew() { - select { - case h.renewed <- struct{}{}: - default: - } +// startHeartbeat takes the send call so tests can count calls instead of reaching a server. +func startHeartbeat(ctx context.Context, send func(context.Context) error) *Heartbeat { + ctx, cancel := context.WithCancel(ctx) + h := &Heartbeat{renewed: make(chan struct{}, 1), cancel: cancel} + go h.run(ctx, send) + return h } -func (h *Heartbeat) run(ctx context.Context) { - defer close(h.done) - - next := time.NewTimer(heartbeatInterval) - defer next.Stop() +func (h *Heartbeat) run(ctx context.Context, send func(context.Context) error) { + timer := time.NewTimer(heartbeatInterval) for { select { case <-ctx.Done(): return - case <-h.stop: - return case <-h.renewed: - next.Reset(heartbeatInterval) - case <-next.C: - h.beat(ctx) - next.Reset(heartbeatInterval) + case <-timer.C: + // A request cut short by Stop is not a failure. + if err := send(ctx); err != nil && ctx.Err() == nil { + logdiag.LogDiag(ctx, diag.Diagnostic{ + Severity: diag.Warning, + Summary: fmt.Sprintf("Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err), + }) + } } + timer.Reset(heartbeatInterval) } } -// beat sends a heartbeat, retrying quickly a few times before leaving it to the next interval. -func (h *Heartbeat) beat(ctx context.Context) { - err := h.send(ctx) - for _, delay := range heartbeatRetryDelays { - if err == nil || ctx.Err() != nil { - break - } - log.Warnf(ctx, "Failed to send deployment heartbeat, retrying in %s: %s", delay, err) - select { - case <-time.After(delay): - case <-ctx.Done(): - return - case <-h.stop: - return - } - err = h.send(ctx) - } - - // A cancelled run needs no warning. - if err != nil && ctx.Err() == nil { - logdiag.LogDiag(ctx, diag.Diagnostic{ - Severity: diag.Warning, - Summary: fmt.Sprintf("Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err), - }) +// renew records that a successful write updated the version, which the service counts as a +// heartbeat, so the next heartbeat is pushed back a full interval. Non-blocking: one pending +// signal is enough. +func (h *Heartbeat) renew() { + select { + case h.renewed <- struct{}{}: + default: } } -// Stop ends the heartbeat and waits for it. Safe to call twice. +// Stop ends the heartbeat, cutting short one in flight. func (h *Heartbeat) Stop() { - h.stopOnce() + h.cancel() } diff --git a/libs/dms/heartbeat_test.go b/libs/dms/heartbeat_test.go index b7d71c4b2f9..b7de3d8bb1d 100644 --- a/libs/dms/heartbeat_test.go +++ b/libs/dms/heartbeat_test.go @@ -73,31 +73,24 @@ func TestHeartbeatSkippedAfterWrite(t *testing.T) { }) } -func TestHeartbeatRetriesThenWaitsForInterval(t *testing.T) { +func TestHeartbeatFailureWarnsAndRetriesNextInterval(t *testing.T) { synctest.Test(t, func(t *testing.T) { ctx := heartbeatContext(t) f := &fakeSend{err: errors.New("connection refused")} h := startHeartbeat(ctx, f.send) defer h.Stop() - // The beat at 30s and its retries 1s, 2s and 4s apart. - time.Sleep(heartbeatInterval + 8*time.Second) + time.Sleep(heartbeatInterval + time.Second) synctest.Wait() - assert.Equal(t, int32(1+len(heartbeatRetryDelays)), f.calls.Load()) + assert.Equal(t, int32(1), f.calls.Load()) diags := logdiag.FlushCollected(ctx) require.Len(t, diags, 1) assert.Equal(t, "Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in 30s: connection refused", diags[0].Summary) - // Then nothing more until a full interval has passed. - time.Sleep(heartbeatInterval - 2*time.Second) - synctest.Wait() - assert.Equal(t, int32(1+len(heartbeatRetryDelays)), f.calls.Load()) - - // Checked before its first retry, 1s later. - time.Sleep(1500 * time.Millisecond) + time.Sleep(heartbeatInterval) synctest.Wait() - assert.Equal(t, int32(2+len(heartbeatRetryDelays)), f.calls.Load()) + assert.Equal(t, int32(2), f.calls.Load()) }) } @@ -106,7 +99,6 @@ func TestHeartbeatStop(t *testing.T) { f := &fakeSend{} h := startHeartbeat(heartbeatContext(t), f.send) - h.Stop() h.Stop() time.Sleep(10 * time.Minute) synctest.Wait() From b99be930399f78e26325b5d7c55c2d693e1dfc37 Mon Sep 17 00:00:00 2001 From: Sankalp-Mittal Date: Thu, 5 Nov 2026 10:57:50 +0000 Subject: [PATCH 3/4] Update heartbeat comments Co-authored-by: Isaac --- libs/dms/heartbeat.go | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/libs/dms/heartbeat.go b/libs/dms/heartbeat.go index 8a2da222ce6..70bcfcc69f2 100644 --- a/libs/dms/heartbeat.go +++ b/libs/dms/heartbeat.go @@ -11,21 +11,17 @@ import ( "github.com/databricks/databricks-sdk-go/service/bundledeployments" ) -// heartbeatInterval is how long a version may go without an update before a heartbeat is sent. -// The UI shows a version as interrupted once its last update is over 2 minutes old. +// heartbeatInterval is the maximum time between consecutive updates. const heartbeatInterval = 30 * time.Second -// Heartbeat keeps an open deployment version's last update time fresh, which is how the UI tells -// a running deploy from an interrupted one. CreateVersion and every successful write update it -// too, so a heartbeat is sent only after an interval without one. Best effort: a failure can only -// make the deploy show as interrupted, so it is reported and the next interval tries again. +// Heartbeat updates the update_time of an open deployment version. This is also updated by +// CreateVersion and every successful write. type Heartbeat struct { renewed chan struct{} cancel context.CancelFunc } -// StartHeartbeat starts heartbeating the version the caller just created. CreateVersion counts -// as the first update, so the first heartbeat goes out an interval later. +// StartHeartbeat starts the heartbeat. We count CreateVersion as the first heartbeat. func StartHeartbeat(ctx context.Context, deploymentID string, versionNum int) *Heartbeat { name := VersionName(deploymentID, versionNum) return startHeartbeat(ctx, func(ctx context.Context) error { @@ -62,9 +58,7 @@ func (h *Heartbeat) run(ctx context.Context, send func(context.Context) error) { } } -// renew records that a successful write updated the version, which the service counts as a -// heartbeat, so the next heartbeat is pushed back a full interval. Non-blocking: one pending -// signal is enough. +// renew records that a successful write updated the version. func (h *Heartbeat) renew() { select { case h.renewed <- struct{}{}: @@ -72,7 +66,7 @@ func (h *Heartbeat) renew() { } } -// Stop ends the heartbeat, cutting short one in flight. +// Stop ends the heartbeat. func (h *Heartbeat) Stop() { h.cancel() } From 85b6141fd817fe89f7506b244dc29520123794ee Mon Sep 17 00:00:00 2001 From: Sankalp-Mittal Date: Thu, 5 Nov 2026 15:38:22 +0000 Subject: [PATCH 4/4] Split the heartbeat send into separate statements Co-authored-by: Isaac --- libs/dms/heartbeat.go | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/libs/dms/heartbeat.go b/libs/dms/heartbeat.go index 70bcfcc69f2..b2dab0a47e3 100644 --- a/libs/dms/heartbeat.go +++ b/libs/dms/heartbeat.go @@ -46,13 +46,18 @@ func (h *Heartbeat) run(ctx context.Context, send func(context.Context) error) { return case <-h.renewed: case <-timer.C: - // A request cut short by Stop is not a failure. - if err := send(ctx); err != nil && ctx.Err() == nil { - logdiag.LogDiag(ctx, diag.Diagnostic{ - Severity: diag.Warning, - Summary: fmt.Sprintf("Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err), - }) + err := send(ctx) + if err == nil { + break } + // We only get an error if Stop was called + if ctx.Err() != nil { + return + } + logdiag.LogDiag(ctx, diag.Diagnostic{ + Severity: diag.Warning, + Summary: fmt.Sprintf("Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err), + }) } timer.Reset(heartbeatInterval) }