Send heartbeats for open DMS deployment versions

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.

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()
+}

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())
+	})
+}

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()