diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 326c3b8354c..de00e76233a 100644 --- a/bundle/direct/dstate/state.go +++ b/bundle/direct/dstate/state.go @@ -248,6 +248,9 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b db.reset() db.mu.Unlock() + // Stopped first: a heartbeat racing the completion would be rejected for nothing. + buf.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..4c239bcc647 --- /dev/null +++ b/libs/dms/heartbeat.go @@ -0,0 +1,72 @@ +package dms + +import ( + "context" + "time" + + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go/service/bundledeployments" +) + +// heartbeatInterval is the maximum time between consecutive updates. +const heartbeatInterval = 30 * time.Second + +// 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 +} + +// startVersionHeartbeat starts the heartbeat. We count CreateVersion as the first heartbeat. +func startVersionHeartbeat(ctx context.Context, deploymentID string, versionNum int) *Heartbeat { + name := VersionName(deploymentID, versionNum) + return startHeartbeat(ctx, func(ctx context.Context) error { + _, err := cmdctx.WorkspaceClient(ctx).BundleDeployments.Heartbeat(ctx, bundledeployments.HeartbeatRequest{Name: name}) + return err + }) +} + +// 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, send func(context.Context) error) { + timer := time.NewTimer(heartbeatInterval) + for { + select { + case <-ctx.Done(): + return + case <-h.renewed: + case <-timer.C: + err := send(ctx) + if err == nil { + break + } + // We only get an error if Stop was called + if ctx.Err() != nil { + return + } + log.Infof(ctx, "Failed to send deployment heartbeat, so the deployment may show as interrupted; retrying in %s: %s", heartbeatInterval, err) + } + timer.Reset(heartbeatInterval) + } +} + +// renew records that a write started, so the next heartbeat waits a full interval. +func (h *Heartbeat) renew() { + select { + case h.renewed <- struct{}{}: + default: + } +} + +// Stop ends the heartbeat. +func (h *Heartbeat) Stop() { + h.cancel() +} diff --git a/libs/dms/heartbeat_test.go b/libs/dms/heartbeat_test.go new file mode 100644 index 00000000000..3b6dc320567 --- /dev/null +++ b/libs/dms/heartbeat_test.go @@ -0,0 +1,94 @@ +package dms + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/stretchr/testify/assert" +) + +// 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 TestHeartbeatBeatsWhileIdle(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + f := &fakeSend{} + h := startHeartbeat(t.Context(), 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(t.Context(), 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 TestHeartbeatFailureRetriesNextInterval(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + f := &fakeSend{err: errors.New("connection refused")} + h := startHeartbeat(t.Context(), f.send) + defer h.Stop() + + time.Sleep(heartbeatInterval + 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 TestHeartbeatStop(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + f := &fakeSend{} + h := startHeartbeat(t.Context(), f.send) + + 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 921651d4325..9c13f543908 100644 --- a/libs/dms/operation_buffer.go +++ b/libs/dms/operation_buffer.go @@ -62,6 +62,10 @@ type OperationBuffer struct { deploymentID string versionNum int + // heartbeat keeps the version's update_time fresh until Stop, past Drain. Each write renews + // it as it starts, since the service counts a write 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 @@ -103,6 +107,8 @@ 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 { + // Started on ctx, not the writer's context below, so it keeps running after Drain. + heartbeat := startVersionHeartbeat(ctx, deploymentID, versionNum) // History can stop independently without canceling the deployment. ctx, cancel := context.WithCancel(ctx) capacity := bufferedOperations @@ -115,6 +121,7 @@ func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum i b := &OperationBuffer{ deploymentID: deploymentID, versionNum: versionNum, + heartbeat: heartbeat, queue: make(chan string, capacity), done: make(chan struct{}), cancel: cancel, @@ -235,6 +242,10 @@ 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, entry *bufferedUpdate) error { + // Renewed when the write starts, not when it succeeds, so a heartbeat only overlaps a write + // that has been in flight for a full interval. Both update the version row. + b.heartbeat.renew() + w := cmdctx.WorkspaceClient(ctx) result, err := w.BundleDeployments.UpdateOperation(ctx, bundledeployments.UpdateOperationRequest{ Name: OperationName(b.deploymentID, b.versionNum, key), @@ -295,6 +306,11 @@ func (b *OperationBuffer) Drain(ctx context.Context) error { } } +// Stop ends the heartbeat, the buffer's only background work left after Drain. +func (b *OperationBuffer) Stop() { + b.heartbeat.Stop() +} + // Stats returns a snapshot of the history measurements. func (b *OperationBuffer) Stats() OperationBufferStats { b.mu.Lock()