diff --git a/bundle/direct/dstate/state.go b/bundle/direct/dstate/state.go index 384246913f4..94f85f4da67 100644 --- a/bundle/direct/dstate/state.go +++ b/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. diff --git a/libs/dms/heartbeat.go b/libs/dms/heartbeat.go new file mode 100644 index 00000000000..b2dab0a47e3 --- /dev/null +++ b/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() +} diff --git a/libs/dms/heartbeat_test.go b/libs/dms/heartbeat_test.go new file mode 100644 index 00000000000..b7de3d8bb1d --- /dev/null +++ b/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()) + }) +} diff --git a/libs/dms/operation_buffer.go b/libs/dms/operation_buffer.go index 7a77cdf83c8..9c62691e5e5 100644 --- a/libs/dms/operation_buffer.go +++ b/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 } diff --git a/libs/dms/operation_buffer_test.go b/libs/dms/operation_buffer_test.go index 0e019284360..546933abc0d 100644 --- a/libs/dms/operation_buffer_test.go +++ b/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