Send heartbeats for open DMS deployment versions

bundle/direct/dstate/state.go

@@ -107,6 +107,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
@@ -187,7 +190,8 @@ func (db *DeploymentState) InitializeOperationBuffer(ctx context.Context, deploy
 		BestEffort:     db.mode == WritingDeploymentHistoryShadow,
 		OperationCount: operationCount,
 	}
-	db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options)
+	db.heartbeat = dms.StartHeartbeat(ctx, deploymentID, versionID)
+	db.operationBuffer = dms.StartOperationBuffer(ctx, deploymentID, versionID, options, db.heartbeat)
 	db.DeploymentID = deploymentID
 	db.VersionID = versionID
 }
@@ -233,9 +237,13 @@ func (db *DeploymentState) CompleteVersion(ctx context.Context, success bool) (b
 	}
 	db.versionCompleted = true
 	deploymentID, versionID := db.DeploymentID, db.VersionID
+	heartbeat := db.heartbeat
 	db.reset()
 	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.

libs/dms/heartbeat.go

@@ -0,0 +1,77 @@
+package dms
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/databricks/cli/libs/cmdctx"
+	"github.com/databricks/cli/libs/diag"
+	"github.com/databricks/cli/libs/logdiag"
+	"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
+}
+
+// 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 {
+		_, 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
+			}
+			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)
+	}
+}
+
+// renew records that a successful write updated the version.
+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,107 @@
+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 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()
+
+		time.Sleep(heartbeatInterval + time.Second)
+		synctest.Wait()
+		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)
+
+		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(heartbeatContext(t), f.send)
+
+		h.Stop()
+		time.Sleep(10 * time.Minute)
+		synctest.Wait()
+		assert.Equal(t, int32(0), f.calls.Load())
+	})
+}

libs/dms/operation_buffer.go

@@ -32,6 +32,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
@@ -69,7 +72,7 @@ type OperationBuffer struct {
 
 // StartOperationBuffer opens the buffer for the version the caller just created. The version
 // must already exist: operations record under it, and nothing here creates it.
-func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum int, options BufferOptions) *OperationBuffer {
+func StartOperationBuffer(ctx context.Context, deploymentID string, versionNum int, options BufferOptions, heartbeat *Heartbeat) *OperationBuffer {
 	// History can stop independently without canceling the deployment.
 	ctx, cancel := context.WithCancel(ctx)
 	capacity := bufferedOperations
@@ -82,6 +85,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,
@@ -203,6 +207,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
 }
 

libs/dms/operation_buffer_test.go

@@ -30,6 +30,11 @@ func queued() *OperationBuffer {
 	}
 }
 
+// idleHeartbeat is a heartbeat whose sends do nothing, for buffers that never write.
+func idleHeartbeat(ctx context.Context) *Heartbeat {
+	return startHeartbeat(ctx, func(context.Context) error { return nil })
+}
+
 func stateUpdate(t *testing.T, name string) OperationUpdate {
 	t.Helper()
 	update, err := NewStateUpdate("id-1", json.RawMessage(`{"state":{"name":"`+name+`"}}`), false)
@@ -74,7 +79,7 @@ func TestBufferFailsOnOversizedState(t *testing.T) {
 
 func TestDrainIsIdempotent(t *testing.T) {
 	// An idle writer must finish when draining begins.
-	b := StartOperationBuffer(t.Context(), "dep-1", 1, BufferOptions{})
+	b := StartOperationBuffer(t.Context(), "dep-1", 1, BufferOptions{}, idleHeartbeat(t.Context()))
 	require.NoError(t, b.Drain(t.Context()))
 	require.NoError(t, b.Drain(t.Context()))
 }
@@ -119,7 +124,7 @@ func TestHistoryQueueDoesNotBlockWhenWriterStops(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})
+		b := StartOperationBuffer(ctx, "dep-1", 1, BufferOptions{BestEffort: true, OperationCount: operationCount}, idleHeartbeat(ctx))
 		synctest.Wait()
 		cancel()
 		<-b.done