Apply stopchan to pgx pollPoolStats

DataDog · Dec 16, 2024 · 2cb3519 · 2cb3519
1 parent 41c3ad9
commit 2cb3519
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 17 deletions.
diff --git a/contrib/jackc/pgx.v5/metrics.go b/contrib/jackc/pgx.v5/metrics.go
@@ -34,23 +34,30 @@ const (
 var interval = 10 * time.Second
 
 // pollPoolStats calls (*pgxpool).Stats on the pool at a predetermined interval. It pushes the pool Stats off to the statsd client.
-func pollPoolStats(statsd internal.StatsdClient, pool *pgxpool.Pool) {
+func pollPoolStats(statsd internal.StatsdClient, pool *pgxpool.Pool, stop chan struct{}) {
 	log.Debug("contrib/jackc/pgx.v5: Traced pool connection found: Pool stats will be gathered and sent every %v.", interval)
-	for range time.NewTicker(interval).C {
-		log.Debug("contrib/jackc/pgx.v5: Reporting pgxpool.Stat metrics...")
-		stat := pool.Stat()
-		statsd.Gauge(AcquireCount, float64(stat.AcquireCount()), []string{}, 1)
-		statsd.Timing(AcquireDuration, stat.AcquireDuration(), []string{}, 1)
-		statsd.Gauge(AcquiredConns, float64(stat.AcquiredConns()), []string{}, 1)
-		statsd.Gauge(CanceledAcquireCount, float64(stat.CanceledAcquireCount()), []string{}, 1)
-		statsd.Gauge(ConstructingConns, float64(stat.ConstructingConns()), []string{}, 1)
-		statsd.Gauge(EmptyAcquireCount, float64(stat.EmptyAcquireCount()), []string{}, 1)
-		statsd.Gauge(IdleConns, float64(stat.IdleConns()), []string{}, 1)
-		statsd.Gauge(MaxConns, float64(stat.MaxConns()), []string{}, 1)
-		statsd.Gauge(TotalConns, float64(stat.TotalConns()), []string{}, 1)
-		statsd.Gauge(NewConnsCount, float64(stat.NewConnsCount()), []string{}, 1)
-		statsd.Gauge(MaxLifetimeDestroyCount, float64(stat.MaxLifetimeDestroyCount()), []string{}, 1)
-		statsd.Gauge(MaxIdleDestroyCount, float64(stat.MaxIdleDestroyCount()), []string{}, 1)
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			log.Debug("contrib/jackc/pgx.v5: Reporting pgxpool.Stat metrics...")
+			stat := pool.Stat()
+			statsd.Gauge(AcquireCount, float64(stat.AcquireCount()), []string{}, 1)
+			statsd.Timing(AcquireDuration, stat.AcquireDuration(), []string{}, 1)
+			statsd.Gauge(AcquiredConns, float64(stat.AcquiredConns()), []string{}, 1)
+			statsd.Gauge(CanceledAcquireCount, float64(stat.CanceledAcquireCount()), []string{}, 1)
+			statsd.Gauge(ConstructingConns, float64(stat.ConstructingConns()), []string{}, 1)
+			statsd.Gauge(EmptyAcquireCount, float64(stat.EmptyAcquireCount()), []string{}, 1)
+			statsd.Gauge(IdleConns, float64(stat.IdleConns()), []string{}, 1)
+			statsd.Gauge(MaxConns, float64(stat.MaxConns()), []string{}, 1)
+			statsd.Gauge(TotalConns, float64(stat.TotalConns()), []string{}, 1)
+			statsd.Gauge(NewConnsCount, float64(stat.NewConnsCount()), []string{}, 1)
+			statsd.Gauge(MaxLifetimeDestroyCount, float64(stat.MaxLifetimeDestroyCount()), []string{}, 1)
+			statsd.Gauge(MaxIdleDestroyCount, float64(stat.MaxIdleDestroyCount()), []string{}, 1)
+		case <-stop:
+			return
+		}
 	}
 }
 

diff --git a/contrib/jackc/pgx.v5/pgxpool.go b/contrib/jackc/pgx.v5/pgxpool.go
@@ -9,6 +9,7 @@ import (
 	"context"
 
 	"github.com/jackc/pgx/v5/pgxpool"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/contribroutines"
 )
 
 func NewPool(ctx context.Context, connString string, opts ...Option) (*pgxpool.Pool, error) {
@@ -30,7 +31,7 @@ func NewPoolWithConfig(ctx context.Context, config *pgxpool.Config, opts ...Opti
 		return nil, err
 	}
 	if tracer.cfg.poolStats && tracer.cfg.statsdClient != nil {
-		go pollPoolStats(tracer.cfg.statsdClient, pool)
+		go pollPoolStats(tracer.cfg.statsdClient, pool, contribroutines.GetStopChan())
 	}
 	return pool, nil
 }