From 6f26c4cce0a1be1885b232bf25c1dc74e3dd7560 Mon Sep 17 00:00:00 2001 From: Ben <52796471+bengentree@users.noreply.github.com> Date: Wed, 10 Jul 2024 04:25:32 -0600 Subject: [PATCH] contrib/jackc/pgx.v5: publish pgxpool stats to statsd client (#2692) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Mikayla Toffler <46911781+mtoffl01@users.noreply.github.com> Co-authored-by: Dario Castañé --- contrib/jackc/pgx.v5/metrics.go | 63 ++++++++++++++++++++++++++++ contrib/jackc/pgx.v5/option.go | 33 ++++++++++++++- contrib/jackc/pgx.v5/option_test.go | 24 +++++++++++ contrib/jackc/pgx.v5/pgx_tracer.go | 1 + contrib/jackc/pgx.v5/pgxpool.go | 12 +++++- contrib/jackc/pgx.v5/pgxpool_test.go | 33 +++++++++++++++ 6 files changed, 163 insertions(+), 3 deletions(-) create mode 100644 contrib/jackc/pgx.v5/metrics.go create mode 100644 contrib/jackc/pgx.v5/option_test.go diff --git a/contrib/jackc/pgx.v5/metrics.go b/contrib/jackc/pgx.v5/metrics.go new file mode 100644 index 0000000000..eb94c50bfc --- /dev/null +++ b/contrib/jackc/pgx.v5/metrics.go @@ -0,0 +1,63 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2022 Datadog, Inc. + +package pgx + +import ( + "time" + + "github.com/jackc/pgx/v5/pgxpool" + "gopkg.in/DataDog/dd-trace-go.v1/internal" + "gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig" + "gopkg.in/DataDog/dd-trace-go.v1/internal/log" +) + +const tracerPrefix = "datadog.tracer." + +const ( + AcquireCount = tracerPrefix + "pgx.pool.connections.acquire" + AcquireDuration = tracerPrefix + "pgx.pool.connections.acquire_duration" + AcquiredConns = tracerPrefix + "pgx.pool.connections.acquired_conns" + CanceledAcquireCount = tracerPrefix + "pgx.pool.connections.canceled_acquire" + ConstructingConns = tracerPrefix + "pgx.pool.connections.constructing_conns" + EmptyAcquireCount = tracerPrefix + "pgx.pool.connections.empty_acquire" + IdleConns = tracerPrefix + "pgx.pool.connections.idle_conns" + MaxConns = tracerPrefix + "pgx.pool.connections.max_conns" + TotalConns = tracerPrefix + "pgx.pool.connections.total_conns" + NewConnsCount = tracerPrefix + "pgx.pool.connections.new_conns" + MaxLifetimeDestroyCount = tracerPrefix + "pgx.pool.connections.max_lifetime_destroy" + MaxIdleDestroyCount = tracerPrefix + "pgx.pool.connections.max_idle_destroy" +) + +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) { + 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) + } +} + +func statsTags(c *config) []string { + tags := globalconfig.StatsTags() + if c.serviceName != "" { + tags = append(tags, "service:"+c.serviceName) + } + return tags +} diff --git a/contrib/jackc/pgx.v5/option.go b/contrib/jackc/pgx.v5/option.go index 83f2ea92d5..b89a3968a0 100644 --- a/contrib/jackc/pgx.v5/option.go +++ b/contrib/jackc/pgx.v5/option.go @@ -5,7 +5,12 @@ package pgx -import "gopkg.in/DataDog/dd-trace-go.v1/internal/namingschema" +import ( + "gopkg.in/DataDog/dd-trace-go.v1/internal" + "gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig" + "gopkg.in/DataDog/dd-trace-go.v1/internal/log" + "gopkg.in/DataDog/dd-trace-go.v1/internal/namingschema" +) type config struct { serviceName string @@ -15,6 +20,8 @@ type config struct { tracePrepare bool traceConnect bool traceAcquire bool + poolStats bool + statsdClient internal.StatsdClient } func defaultConfig() *config { @@ -29,6 +36,21 @@ func defaultConfig() *config { } } +// checkStatsdRequired adds a statsdClient onto the config if poolStats is enabled +// NOTE: For now, the only use-case for a statsdclient is the poolStats feature. If a statsdclient becomes necessary for other items in future work, then this logic should change +func (c *config) checkStatsdRequired() { + if c.poolStats && c.statsdClient == nil { + // contrib/jackc/pgx's statsdclient should always inherit its address from the tracer's statsdclient via the globalconfig + // destination is not user-configurable + sc, err := internal.NewStatsdClient(globalconfig.DogstatsdAddr(), statsTags(c)) + if err == nil { + c.statsdClient = sc + } else { + log.Warn("contrib/jackc/pgx.v5: Error creating statsd client; Pool stats will be dropped: %v", err) + } + } +} + type Option func(*config) // WithServiceName sets the service name to use for all spans. @@ -90,3 +112,12 @@ func WithTraceConnect(enabled bool) Option { c.traceConnect = enabled } } + +// WithPoolStats enables polling of pgxpool.Stat metrics +// ref: https://pkg.go.dev/github.com/jackc/pgx/v5/pgxpool#Stat +// These metrics are submitted to Datadog and are not billed as custom metrics +func WithPoolStats() Option { + return func(cfg *config) { + cfg.poolStats = true + } +} diff --git a/contrib/jackc/pgx.v5/option_test.go b/contrib/jackc/pgx.v5/option_test.go new file mode 100644 index 0000000000..57741c1aa8 --- /dev/null +++ b/contrib/jackc/pgx.v5/option_test.go @@ -0,0 +1,24 @@ +// Unless explicitly stated otherwise all files in this repository are licensed +// under the Apache License Version 2.0. +// This product includes software developed at Datadog (https://www.datadoghq.com/). +// Copyright 2022 Datadog, Inc. + +package pgx + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestWithPoolStats(t *testing.T) { + t.Run("default off", func(t *testing.T) { + cfg := defaultConfig() + assert.False(t, cfg.poolStats) + }) + t.Run("on", func(t *testing.T) { + cfg := new(config) + WithPoolStats()(cfg) + assert.True(t, cfg.poolStats) + }) +} diff --git a/contrib/jackc/pgx.v5/pgx_tracer.go b/contrib/jackc/pgx.v5/pgx_tracer.go index 77249f2248..f361a3ac48 100644 --- a/contrib/jackc/pgx.v5/pgx_tracer.go +++ b/contrib/jackc/pgx.v5/pgx_tracer.go @@ -63,6 +63,7 @@ func newPgxTracer(opts ...Option) *pgxTracer { for _, opt := range opts { opt(cfg) } + cfg.checkStatsdRequired() return &pgxTracer{cfg: cfg} } diff --git a/contrib/jackc/pgx.v5/pgxpool.go b/contrib/jackc/pgx.v5/pgxpool.go index f61a162948..f4a8cb7ed0 100644 --- a/contrib/jackc/pgx.v5/pgxpool.go +++ b/contrib/jackc/pgx.v5/pgxpool.go @@ -20,6 +20,14 @@ func NewPool(ctx context.Context, connString string, opts ...Option) (*pgxpool.P } func NewPoolWithConfig(ctx context.Context, config *pgxpool.Config, opts ...Option) (*pgxpool.Pool, error) { - config.ConnConfig.Tracer = newPgxTracer(opts...) - return pgxpool.NewWithConfig(ctx, config) + tracer := newPgxTracer(opts...) + config.ConnConfig.Tracer = tracer + pool, err := pgxpool.NewWithConfig(ctx, config) + if err != nil { + return nil, err + } + if tracer.cfg.poolStats && tracer.cfg.statsdClient != nil { + go pollPoolStats(tracer.cfg.statsdClient, pool) + } + return pool, nil } diff --git a/contrib/jackc/pgx.v5/pgxpool_test.go b/contrib/jackc/pgx.v5/pgxpool_test.go index 80ab40e79f..c784da33a3 100644 --- a/contrib/jackc/pgx.v5/pgxpool_test.go +++ b/contrib/jackc/pgx.v5/pgxpool_test.go @@ -8,8 +8,11 @@ package pgx import ( "context" "testing" + "time" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer" + "gopkg.in/DataDog/dd-trace-go.v1/internal" + "gopkg.in/DataDog/dd-trace-go.v1/internal/statsdtest" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -38,3 +41,33 @@ func TestPool(t *testing.T) { assert.Len(t, mt.OpenSpans(), 0) assert.Len(t, mt.FinishedSpans(), 7) } + +func TestPoolWithPoolStats(t *testing.T) { + originalInterval := interval + interval = 1 * time.Millisecond + t.Cleanup(func() { + interval = originalInterval + }) + + ctx := context.Background() + statsd := new(statsdtest.TestStatsdClient) + conn, err := NewPool(ctx, postgresDSN, withStatsdClient(statsd), WithPoolStats()) + require.NoError(t, err) + defer conn.Close() + + wantStats := []string{AcquireCount, AcquireDuration, AcquiredConns, CanceledAcquireCount, ConstructingConns, EmptyAcquireCount, IdleConns, MaxConns, TotalConns, NewConnsCount, MaxLifetimeDestroyCount, MaxIdleDestroyCount} + + assert := assert.New(t) + if err := statsd.Wait(assert, len(wantStats), time.Second); err != nil { + t.Fatalf("statsd.Wait(): %v", err) + } + for _, name := range wantStats { + assert.Contains(statsd.CallNames(), name) + } +} + +func withStatsdClient(s internal.StatsdClient) Option { + return func(c *config) { + c.statsdClient = s + } +}