Merge pull request #714 from ay901246/task-heartbeat-logger

Alphasite · web-flow · commit 9e93907ea3b7 · 2026-03-16T18:26:46.000-07:00
Add --with-heartbeat flag for heartbeat logging during long-running tasks
diff --git a/cmd/cmd.go b/cmd/cmd.go
@@ -224,7 +224,19 @@ func (c Cmd) Execute() (cmdErr error) {
 		return NewErrandsCmd(deps.UI, c.deployment()).Run()
 
 	case *RunErrandOpts:
-		director, deployment := c.directorAndDeployment()
+		sess, ok := c.session().(*SessionImpl)
+		if !ok {
+			return fmt.Errorf("internal error: expected *SessionImpl")
+		}
+		director, err := sess.Director()
+		c.panicIfErr(err)
+		deployment, err := sess.Deployment()
+		c.panicIfErr(err)
+
+		if opts.WithHeartbeat != nil && sess.taskReporter != nil {
+			sess.taskReporter.EnableWithHeartbeat(time.Duration(*opts.WithHeartbeat) * time.Second)
+		}
+
 		downloader := NewUIDownloader(director, deps.Time, deps.FS, deps.UI)
 		return NewRunErrandCmd(deployment, downloader, deps.UI).Run(*opts)
 
@@ -546,7 +558,10 @@ func (c Cmd) config() cmdconf.Config {
 }
 
 func (c Cmd) session() Session {
-	return NewSessionFromOpts(c.BoshOpts, c.config(), c.deps.UI, true, true, c.deps.FS, c.deps.Logger)
+	return NewSessionImpl(
+		NewSessionContextImpl(c.BoshOpts, c.config(), c.deps.FS),
+		c.deps.UI, true, true, c.deps.Logger,
+	)
 }
 
 func (c Cmd) director() boshdir.Director {
diff --git a/cmd/opts/opts.go b/cmd/opts/opts.go
@@ -713,6 +713,8 @@ type RunErrandOpts struct {
 	KeepAlive   bool `long:"keep-alive" description:"Use existing VM to run an errand and keep it after completion"`
 	WhenChanged bool `long:"when-changed" description:"Run errand only if errand configuration has changed or if the previous run was unsuccessful"`
 
+	WithHeartbeat *int `long:"with-heartbeat" description:"Print task state every N seconds while waiting. Use '=' to specify interval" optional:"true" optional-value:"30"`
+
 	DownloadLogs  bool        `long:"download-logs" description:"Download logs"`
 	LogsDirectory DirOrCWDArg `long:"logs-dir" description:"Destination directory for logs" default:"."`
 
diff --git a/cmd/opts/opts_test.go b/cmd/opts/opts_test.go
@@ -2150,6 +2150,14 @@ var _ = Describe("Opts", func() {
 			})
 		})
 
+		Describe("WithHeartbeat", func() {
+			It("contains desired values", func() {
+				Expect(getStructTagForName("WithHeartbeat", opts)).To(Equal(
+					`long:"with-heartbeat" description:"Print task state every N seconds while waiting. Use '=' to specify interval" optional:"true" optional-value:"30"`,
+				))
+			})
+		})
+
 		Describe("DownloadLogs", func() {
 			It("contains desired values", func() {
 				Expect(getStructTagForName("DownloadLogs", opts)).To(Equal(
diff --git a/cmd/session.go b/cmd/session.go
@@ -22,6 +22,7 @@ type SessionImpl struct {
 
 	// Memoized
 	director        boshdir.Director
+	taskReporter    *boshuit.ReporterImpl
 	directorInfo    boshdir.Info
 	directorInfoSet bool
 }
@@ -118,10 +119,10 @@ func (c *SessionImpl) Director() (boshdir.Director, error) {
 		c.ui.PrintLinef("Using environment '%s' as %s", c.Environment(), creds.Description())
 	}
 
-	taskReporter := boshuit.NewReporter(c.ui, true)
+	c.taskReporter = boshuit.NewReporter(c.ui, true)
 	fileReporter := boshui.NewFileReporter(c.ui)
 
-	director, err := boshdir.NewFactory(c.logger).New(dirConfig, taskReporter, fileReporter)
+	director, err := boshdir.NewFactory(c.logger).New(dirConfig, c.taskReporter, fileReporter)
 	if err != nil {
 		return nil, err
 	}
diff --git a/director/directorfakes/fake_task_reporter.go b/director/directorfakes/fake_task_reporter.go
diff --git a/director/interfaces.go b/director/interfaces.go
@@ -304,6 +304,7 @@ type TaskReporter interface {
 	TaskStarted(int)
 	TaskFinished(int, string)
 	TaskOutputChunk(int, []byte)
+	TaskHeartbeat(id int, state string, startedAt int64)
 }
 
 //counterfeiter:generate . OrphanDisk
diff --git a/director/noop_reporters.go b/director/noop_reporters.go
@@ -43,6 +43,7 @@ func NewNoopTaskReporter() NoopTaskReporter {
 	return NoopTaskReporter{}
 }
 
-func (r NoopTaskReporter) TaskStarted(id int)                   {}
-func (r NoopTaskReporter) TaskFinished(id int, state string)    {}
-func (r NoopTaskReporter) TaskOutputChunk(id int, chunk []byte) {}
+func (r NoopTaskReporter) TaskStarted(id int)                                  {}
+func (r NoopTaskReporter) TaskFinished(id int, state string)                   {}
+func (r NoopTaskReporter) TaskOutputChunk(id int, chunk []byte)                {}
+func (r NoopTaskReporter) TaskHeartbeat(id int, state string, startedAt int64) {}
diff --git a/director/task_client_request.go b/director/task_client_request.go
@@ -27,8 +27,9 @@ func NewTaskClientRequest(
 }
 
 type taskShortResp struct {
-	ID    int    // 165
-	State string // e.g. "queued", "processing", "done", "error", "cancelled"
+	ID        int    `json:"id"`
+	State     string `json:"state"`      // e.g. "queued", "processing", "done", "error", "cancelled"
+	StartedAt int64  `json:"started_at"` // 1440318199
 }
 
 func (r taskShortResp) IsRunning() bool {
@@ -111,6 +112,7 @@ func (r TaskClientRequest) WaitForCompletion(id int, type_ string, taskReporter
 		}
 
 		if taskResp.IsRunning() {
+			taskReporter.TaskHeartbeat(taskResp.ID, taskResp.State, taskResp.StartedAt)
 			time.Sleep(r.taskCheckStepDuration)
 			continue
 		}
diff --git a/director/task_client_request_test.go b/director/task_client_request_test.go
@@ -246,6 +246,61 @@ var _ = Describe("TaskClientRequest", func() {
 		})
 	})
 
+	Describe("WaitForCompletion heartbeat", func() {
+		It("emits a heartbeat for a processing task", func() {
+			hbReporter := &fakedir.FakeTaskReporter{}
+			hbReq := buildReq(hbReporter)
+
+			server.AppendHandlers(
+				ghttp.CombineHandlers(
+					ghttp.VerifyRequest("GET", "/tasks/42"),
+					ghttp.RespondWith(http.StatusOK, `{"id":42,"state":"processing","description":"run errand 'smoke'","started_at":1700000000}`),
+				),
+				ghttp.CombineHandlers(
+					ghttp.VerifyRequest("GET", "/tasks/42/output", "type=event"),
+					ghttp.RespondWith(http.StatusOK, ""),
+				),
+				ghttp.CombineHandlers(
+					ghttp.VerifyRequest("GET", "/tasks/42"),
+					ghttp.RespondWith(http.StatusOK, `{"id":42,"state":"done"}`),
+				),
+				ghttp.CombineHandlers(
+					ghttp.VerifyRequest("GET", "/tasks/42/output", "type=event"),
+					ghttp.RespondWith(http.StatusOK, ""),
+				),
+			)
+
+			err := hbReq.WaitForCompletion(42, "event", hbReporter)
+			Expect(err).ToNot(HaveOccurred())
+
+			Expect(hbReporter.TaskHeartbeatCallCount()).To(BeNumerically(">=", 1))
+			id, state, startedAt := hbReporter.TaskHeartbeatArgsForCall(0)
+			Expect(id).To(Equal(42))
+			Expect(state).To(Equal("processing"))
+			Expect(startedAt).To(Equal(int64(1700000000)))
+		})
+
+		It("does not emit heartbeats for tasks that immediately finish", func() {
+			hbReporter := &fakedir.FakeTaskReporter{}
+			hbReq := buildReq(hbReporter)
+
+			server.AppendHandlers(
+				ghttp.CombineHandlers(
+					ghttp.VerifyRequest("GET", "/tasks/42"),
+					ghttp.RespondWith(http.StatusOK, `{"id":42,"state":"done"}`),
+				),
+				ghttp.CombineHandlers(
+					ghttp.VerifyRequest("GET", "/tasks/42/output", "type=event"),
+					ghttp.RespondWith(http.StatusOK, ""),
+				),
+			)
+
+			err := hbReq.WaitForCompletion(42, "event", hbReporter)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(hbReporter.TaskHeartbeatCallCount()).To(Equal(0))
+		})
+	})
+
 	Describe("WaitForCompletion", func() {
 		var (
 			taskReporter *fakedir.FakeTaskReporter
diff --git a/ui/task/interfaces.go b/ui/task/interfaces.go
@@ -4,6 +4,7 @@ type Reporter interface {
 	TaskStarted(int)
 	TaskFinished(int, string)
 	TaskOutputChunk(int, []byte)
+	TaskHeartbeat(id int, state string, startedAt int64)
 }
 
 type Task interface {
diff --git a/ui/task/reporter.go b/ui/task/reporter.go
@@ -5,17 +5,22 @@ import (
 	"fmt"
 	"strings"
 	"sync"
+	"time"
 
 	"golang.org/x/text/cases"
 	"golang.org/x/text/language"
 
 	boshui "github.com/cloudfoundry/bosh-cli/v7/ui"
+	boshuifmt "github.com/cloudfoundry/bosh-cli/v7/ui/fmt"
 )
 
 type ReporterImpl struct {
 	ui          boshui.UI
 	isForEvents bool
 
+	withHeartbeatInterval time.Duration
+	lastHeartbeat         time.Time
+
 	events          map[int][]*Event
 	eventMarkers    []eventMarker
 	lastGlobalEvent *Event
@@ -45,6 +50,12 @@ func NewReporter(ui boshui.UI, isForEvents bool) *ReporterImpl {
 	}
 }
 
+func (r *ReporterImpl) EnableWithHeartbeat(interval time.Duration) {
+	r.Lock()
+	defer r.Unlock()
+	r.withHeartbeatInterval = interval
+}
+
 func (r *ReporterImpl) TaskStarted(id int) {
 	r.Lock()
 	defer r.Unlock()
@@ -111,6 +122,30 @@ func (r *ReporterImpl) TaskOutputChunk(id int, chunk []byte) {
 	r.eventMarkers = append(r.eventMarkers, eventMarker{TaskID: id, Type: taskOutput})
 }
 
+func (r *ReporterImpl) TaskHeartbeat(id int, state string, startedAt int64) {
+	r.Lock()
+	defer r.Unlock()
+
+	if r.withHeartbeatInterval <= 0 {
+		return
+	}
+
+	now := time.Now()
+	if !r.lastHeartbeat.IsZero() && now.Sub(r.lastHeartbeat) < r.withHeartbeatInterval {
+		return
+	}
+	r.lastHeartbeat = now
+
+	msg := "Task state: " + state
+	if state != "queued" && startedAt > 0 {
+		elapsed := time.Since(time.Unix(startedAt, 0)).Truncate(time.Second)
+		msg += fmt.Sprintf(" (%s elapsed)", elapsed)
+	}
+
+	r.printBlock(fmt.Sprintf("\nTask %d | %s | ", id, now.UTC().Format(boshuifmt.TimeHoursFmt)))
+	r.printBlock(msg)
+}
+
 func (r *ReporterImpl) showEvent(id int, str string) {
 	event := Event{TaskID: id}
 
diff --git a/ui/task/reporter_test.go b/ui/task/reporter_test.go
@@ -2,6 +2,9 @@ package task_test
 
 import (
 	"bytes"
+	"fmt"
+	"strings"
+	"time"
 
 	boshlog "github.com/cloudfoundry/bosh-utils/logger"
 	. "github.com/onsi/ginkgo/v2"
@@ -335,4 +338,53 @@ Task 101 done
 `))
 		})
 	})
+
+	Describe("TaskHeartbeat", func() {
+		It("does not print when heartbeat is not enabled", func() {
+			impl := boshuit.NewReporter(fakeUI, true)
+			impl.TaskHeartbeat(42, "processing", int64(1700000000))
+			Expect(fakeUI.Blocks).To(BeNil())
+		})
+
+		It("prints state for a queued task when heartbeat is enabled", func() {
+			impl := boshuit.NewReporter(fakeUI, true)
+			impl.EnableWithHeartbeat(10 * time.Second)
+			impl.TaskHeartbeat(42, "queued", int64(0))
+			Expect(len(fakeUI.Blocks)).To(BeNumerically(">=", 1))
+			combined := fmt.Sprintf("%v", fakeUI.Blocks)
+			Expect(combined).To(ContainSubstring("Task state: queued"))
+			Expect(combined).NotTo(ContainSubstring("elapsed"))
+		})
+
+		It("prints state with elapsed time for a processing task with startedAt", func() {
+			impl := boshuit.NewReporter(fakeUI, true)
+			impl.EnableWithHeartbeat(10 * time.Second)
+			impl.TaskHeartbeat(42, "processing", int64(1700000000))
+			Expect(len(fakeUI.Blocks)).To(BeNumerically(">=", 1))
+			combined := fmt.Sprintf("%v", fakeUI.Blocks)
+			Expect(combined).To(ContainSubstring("Task state: processing"))
+			Expect(combined).To(ContainSubstring("elapsed"))
+		})
+
+		It("prints state without elapsed time when processing but startedAt is 0", func() {
+			impl := boshuit.NewReporter(fakeUI, true)
+			impl.EnableWithHeartbeat(10 * time.Second)
+			impl.TaskHeartbeat(42, "processing", int64(0))
+			Expect(len(fakeUI.Blocks)).To(BeNumerically(">=", 1))
+			combined := fmt.Sprintf("%v", fakeUI.Blocks)
+			Expect(combined).To(ContainSubstring("Task state: processing"))
+			Expect(combined).NotTo(ContainSubstring("elapsed"))
+			Expect(combined).NotTo(ContainSubstring("queued"))
+		})
+
+		It("throttles heartbeat output based on the configured interval", func() {
+			impl := boshuit.NewReporter(fakeUI, true)
+			impl.EnableWithHeartbeat(1 * time.Hour)
+			impl.TaskHeartbeat(42, "processing", int64(1700000000))
+			impl.TaskHeartbeat(42, "processing", int64(1700000000))
+			impl.TaskHeartbeat(42, "processing", int64(1700000000))
+			combined := fmt.Sprintf("%v", fakeUI.Blocks)
+			Expect(strings.Count(combined, "Task state")).To(Equal(1))
+		})
+	})
 })
diff --git a/ui/task/taskfakes/fake_reporter.go b/ui/task/taskfakes/fake_reporter.go

Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,7 @@ type SessionImpl struct {`
`22`	`22`
`23`	`23`	`// Memoized`
`24`	`24`	`director boshdir.Director`
	`25`	`+ taskReporter *boshuit.ReporterImpl`
`25`	`26`	`directorInfo boshdir.Info`
`26`	`27`	`directorInfoSet bool`
`27`	`28`	`}`
`@@ -118,10 +119,10 @@ func (c *SessionImpl) Director() (boshdir.Director, error) {`
`118`	`119`	`c.ui.PrintLinef("Using environment '%s' as %s", c.Environment(), creds.Description())`
`119`	`120`	`}`
`120`	`121`
`121`		`- taskReporter := boshuit.NewReporter(c.ui, true)`
	`122`	`+ c.taskReporter = boshuit.NewReporter(c.ui, true)`
`122`	`123`	`fileReporter := boshui.NewFileReporter(c.ui)`
`123`	`124`
`124`		`- director, err := boshdir.NewFactory(c.logger).New(dirConfig, taskReporter, fileReporter)`
	`125`	`+ director, err := boshdir.NewFactory(c.logger).New(dirConfig, c.taskReporter, fileReporter)`
`125`	`126`	`if err != nil {`
`126`	`127`	`return nil, err`
`127`	`128`	`}`
Original file line number	Diff line number	Diff line change
`@@ -304,6 +304,7 @@ type TaskReporter interface {`
`304`	`304`	`TaskStarted(int)`
`305`	`305`	`TaskFinished(int, string)`
`306`	`306`	`TaskOutputChunk(int, []byte)`
	`307`	`+ TaskHeartbeat(id int, state string, startedAt int64)`
`307`	`308`	`}`
`308`	`309`
`309`	`310`	`//counterfeiter:generate . OrphanDisk`
Original file line number	Diff line number	Diff line change
`@@ -27,8 +27,9 @@ func NewTaskClientRequest(`
`27`	`27`	`}`
`28`	`28`
`29`	`29`	`type taskShortResp struct {`
`30`		`- ID int // 165`
`31`		`- State string // e.g. "queued", "processing", "done", "error", "cancelled"`
	`30`	+ ID int `json:"id"`
	`31`	+ State string `json:"state"` // e.g. "queued", "processing", "done", "error", "cancelled"
	`32`	+ StartedAt int64 `json:"started_at"` // 1440318199
`32`	`33`	`}`
`33`	`34`
`34`	`35`	`func (r taskShortResp) IsRunning() bool {`
`@@ -111,6 +112,7 @@ func (r TaskClientRequest) WaitForCompletion(id int, type_ string, taskReporter`
`111`	`112`	`}`
`112`	`113`
`113`	`114`	`if taskResp.IsRunning() {`
	`115`	`+ taskReporter.TaskHeartbeat(taskResp.ID, taskResp.State, taskResp.StartedAt)`
`114`	`116`	`time.Sleep(r.taskCheckStepDuration)`
`115`	`117`	`continue`
`116`	`118`	`}`
Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,7 @@ type Reporter interface {`
`4`	`4`	`TaskStarted(int)`
`5`	`5`	`TaskFinished(int, string)`
`6`	`6`	`TaskOutputChunk(int, []byte)`
	`7`	`+ TaskHeartbeat(id int, state string, startedAt int64)`
`7`	`8`	`}`
`8`	`9`
`9`	`10`	`type Task interface {`