Stop async measurements on pubsub close

dannykopping · dannykopping · commit bf89a2db8e62 · 2024-05-07T10:23:04.000Z
Refactor async measurement for immediate exit upon signal

Signed-off-by: Danny Kopping &lt;danny@coder.com&gt;
diff --git a/coderd/database/pubsub/latency.go b/coderd/database/pubsub/latency.go
@@ -20,8 +20,11 @@ type LatencyMeasurer struct {
 	channel uuid.UUID
 	logger  slog.Logger
 
+	// background measurement members
 	collections atomic.Int64
 	last        atomic.Value
+	asyncTick   *time.Ticker
+	stop        chan struct{}
 }
 
 type LatencyMeasurement struct {
@@ -36,6 +39,7 @@ func NewLatencyMeasurer(logger slog.Logger) *LatencyMeasurer {
 	return &LatencyMeasurer{
 		channel: uuid.New(),
 		logger:  logger,
+		stop:    make(chan struct{}, 1),
 	}
 }
 
@@ -47,6 +51,7 @@ func (lm *LatencyMeasurer) Measure(ctx context.Context, p Pubsub) LatencyMeasure
 	)
 
 	msg := []byte(uuid.New().String())
+	lm.logger.Debug(ctx, "performing measurement", slog.F("msg", msg))
 
 	cancel, err := p.Subscribe(lm.latencyChannelName(), func(ctx context.Context, in []byte) {
 		if !bytes.Equal(in, msg) {
@@ -81,23 +86,32 @@ func (lm *LatencyMeasurer) Measure(ctx context.Context, p Pubsub) LatencyMeasure
 // MeasureAsync runs latency measurements asynchronously on a given interval.
 // This function is expected to be run in a goroutine and will exit when the context is canceled.
 func (lm *LatencyMeasurer) MeasureAsync(ctx context.Context, p Pubsub, interval time.Duration) {
-	tick := time.NewTicker(interval)
-	defer tick.Stop()
-
-	for ; true; <-tick.C { // tick immediately
-		select {
-		case <-ctx.Done():
+	lm.asyncTick = time.NewTicker(interval)
+	defer lm.asyncTick.Stop()
+
+loop:
+	for {
+		// run immediately on first call, then sleep a tick before each invocation
+		if p == nil {
+			lm.logger.Error(ctx, "given pubsub is nil")
 			return
-		default:
-			if p == nil {
-				lm.logger.Error(ctx, "given pubsub is nil")
-				return
-			}
 		}
 
 		lm.collections.Add(1)
 		measure := lm.Measure(ctx, p)
 		lm.last.Store(&measure)
+
+		select {
+		case <-lm.asyncTick.C:
+			continue
+
+		// bail out if signaled
+		case <-lm.stop:
+			break loop
+		case <-ctx.Done():
+			lm.logger.Debug(ctx, "async measurement context canceled", slog.Error(ctx.Err()))
+			break loop
+		}
 	}
 }
 
@@ -115,6 +129,15 @@ func (lm *LatencyMeasurer) MeasurementCount() int64 {
 	return lm.collections.Load()
 }
 
+// Stop stops any background measurements.
+func (lm *LatencyMeasurer) Stop() {
+	if lm.asyncTick == nil {
+		return
+	}
+	lm.asyncTick.Stop()
+	lm.stop <- struct{}{}
+}
+
 func (lm *LatencyMeasurer) latencyChannelName() string {
 	return fmt.Sprintf("latency-measure:%s", lm.channel)
 }
diff --git a/coderd/database/pubsub/pubsub.go b/coderd/database/pubsub/pubsub.go
@@ -311,6 +311,8 @@ func (p *PGPubsub) Close() error {
 	err := p.closeListener()
 	<-p.listenDone
 	p.logger.Debug(context.Background(), "pubsub closed")
+	p.latencyMeasurer.Stop()
+	p.logger.Debug(context.Background(), "background latency measurement has stopped")
 	return err
 }
 

Original file line number	Diff line number	Diff line change
`@@ -311,6 +311,8 @@ func (p *PGPubsub) Close() error {`
`311`	`311`	`err := p.closeListener()`
`312`	`312`	`<-p.listenDone`
`313`	`313`	`p.logger.Debug(context.Background(), "pubsub closed")`
	`314`	`+ p.latencyMeasurer.Stop()`
	`315`	`+ p.logger.Debug(context.Background(), "background latency measurement has stopped")`
`314`	`316`	`return err`
`315`	`317`	`}`
`316`	`318`