Ensure heartbeats run until agent is stopped

Fixes a regression in #971

agent/agent_worker.go

@@ -1,6 +1,7 @@
 package agent
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"sync"
@@ -136,6 +137,10 @@ func (a *AgentWorker) Start() error {
 	// Create the ticker
 	a.ticker = time.NewTicker(pingInterval)
 
+	// Use a context to run heartbeats for as long as the agent runs for
+	heartbeatCtx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	// Setup and start the heartbeater
 	go func() {
 		for {
@@ -150,7 +155,7 @@ func (a *AgentWorker) Start() error {
 						err, heartbeatInterval, time.Now().Sub(lastHeartbeat))
 				}
 
-			case <-a.stop:
+			case <-heartbeatCtx.Done():
 				a.logger.Debug("Stopping heartbeats")
 				return
 			}