Revert "fix(coderd/provisionerdserver): fix test flake in TestHeartbeat"

johnstcn · johnstcn · commit 6449c93eb8e9 · 2024-01-25T11:09:46.000Z
This reverts commit bb32c58.
diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
@@ -72,9 +72,6 @@ type Options struct {
 	// The default function just calls UpdateProvisionerDaemonLastSeenAt.
 	// This is mainly used for testing.
 	HeartbeatFn func(context.Context) error
-
-	// HeartbeatDone is used for testing.
-	HeartbeatDone chan struct{}
 }
 
 type server struct {
@@ -186,9 +183,6 @@ func NewServer(
 	if options.HeartbeatInterval == 0 {
 		options.HeartbeatInterval = DefaultHeartbeatInterval
 	}
-	if options.HeartbeatDone == nil {
-		options.HeartbeatDone = make(chan struct{})
-	}
 
 	s := &server{
 		lifecycleCtx:                lifecycleCtx,
@@ -219,7 +213,7 @@ func NewServer(
 		s.heartbeatFn = s.defaultHeartbeat
 	}
 
-	go s.heartbeatLoop(options.HeartbeatDone)
+	go s.heartbeatLoop()
 	return s, nil
 }
 
@@ -233,21 +227,17 @@ func (s *server) timeNow() time.Time {
 }
 
 // heartbeatLoop runs heartbeatOnce at the interval specified by HeartbeatInterval
-// until the lifecycle context is canceled. Done is closed on exit.
-func (s *server) heartbeatLoop(hbDone chan<- struct{}) {
+// until the lifecycle context is canceled.
+func (s *server) heartbeatLoop() {
 	tick := time.NewTicker(time.Nanosecond)
-	defer func() {
-		close(hbDone)
-	}()
+	defer tick.Stop()
 	for {
 		select {
 		case <-s.lifecycleCtx.Done():
 			s.Logger.Debug(s.lifecycleCtx, "heartbeat loop canceled")
-			tick.Stop()
 			return
 		case <-tick.C:
 			if s.lifecycleCtx.Err() != nil {
-				tick.Stop()
 				return
 			}
 			start := s.timeNow()
diff --git a/coderd/provisionerdserver/provisionerdserver_test.go b/coderd/provisionerdserver/provisionerdserver_test.go
@@ -104,7 +104,6 @@ func TestHeartbeat(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	t.Cleanup(cancel)
 	heartbeatChan := make(chan struct{})
-	heartbeatDone := make(chan struct{})
 	heartbeatFn := func(hbCtx context.Context) error {
 		t.Logf("heartbeat")
 		select {
@@ -118,7 +117,6 @@ func TestHeartbeat(t *testing.T) {
 	//nolint:dogsled // ｡:ﾟ૮ ˶ˆ ﻌ ˆ˶ ა ﾟ:｡
 	_, _, _, _ = setup(t, false, &overrides{
 		ctx:               ctx,
-		heartbeatDone:     heartbeatDone,
 		heartbeatFn:       heartbeatFn,
 		heartbeatInterval: testutil.IntervalFast,
 	})
@@ -127,9 +125,17 @@ func TestHeartbeat(t *testing.T) {
 	require.True(t, ok, "first heartbeat not received")
 	_, ok = <-heartbeatChan
 	require.True(t, ok, "second heartbeat not received")
-	// Cancel the context. This should cause heartbeatDone to be closed.
 	cancel()
-	<-heartbeatDone
+	// Close the channel to ensure we don't receive any more heartbeats.
+	// The test will fail if we do.
+	defer func() {
+		if r := recover(); r != nil {
+			t.Fatalf("heartbeat received after cancel: %v", r)
+		}
+	}()
+
+	close(heartbeatChan)
+	<-time.After(testutil.IntervalMedium)
 }
 
 func TestAcquireJob(t *testing.T) {
@@ -1721,7 +1727,6 @@ type overrides struct {
 	acquireJobLongPollDuration  time.Duration
 	heartbeatFn                 func(ctx context.Context) error
 	heartbeatInterval           time.Duration
-	heartbeatDone               chan struct{}
 }
 
 func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisionerDaemonServer, database.Store, pubsub.Pubsub, database.ProvisionerDaemon) {
@@ -1746,9 +1751,6 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
 	if ov.heartbeatInterval == 0 {
 		ov.heartbeatInterval = testutil.IntervalMedium
 	}
-	if ov.heartbeatDone == nil {
-		ov.heartbeatDone = make(chan struct{})
-	}
 	if ov.deploymentValues != nil {
 		deploymentValues = ov.deploymentValues
 	}
@@ -1813,7 +1815,6 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
 			AcquireJobLongPollDur: pollDur,
 			HeartbeatInterval:     ov.heartbeatInterval,
 			HeartbeatFn:           ov.heartbeatFn,
-			HeartbeatDone:         ov.heartbeatDone,
 		},
 	)
 	require.NoError(t, err)
