contrib/database/sql: Submit DBStats as Datadog metrics

contrib/database/sql/option.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"reflect"
 	"strings"
+	"time"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal"
@@ -29,6 +30,7 @@ type config struct {
 	errCheck           func(err error) bool
 	tags               map[string]interface{}
 	dbmPropagationMode tracer.DBMPropagationMode
+	dbStats            time.Duration
 }
 
 func (c *config) checkDBMPropagation(driverName string, driver driver.Driver, dsn string) {
@@ -262,3 +264,13 @@ func WithDBMPropagation(mode tracer.DBMPropagationMode) Option {
 		cfg.dbmPropagationMode = mode
 	}
 }
+
+// WithDBStats enables polling of DBStats metrics on the registered DB at the specified interval
+// the interval must be non-zero to enable the feature
+// ref: https://pkg.go.dev/database/sql#DBStats
+// These metrics are submitted to Datadog and are not billed as custom metrics
+func WithDBStats(interval time.Duration) Option {
+	return func(cfg *config) {
+		cfg.dbStats = interval
+	}
+}

contrib/database/sql/option_test.go

@@ -7,6 +7,7 @@ package sql
 
 import (
 	"testing"
+	"time"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig"
 
@@ -43,3 +44,27 @@ func TestAnalyticsSettings(t *testing.T) {
 		assert.Equal(t, 0.2, cfg.analyticsRate)
 	})
 }
+
+func TestWithDBStats(t *testing.T) {
+	t.Run("default off", func(t *testing.T) {
+		cfg := new(config)
+		defaults(cfg, "", nil)
+		assert.True(t, int64(cfg.dbStats) == 0)
+		assert.False(t, dbStatsEnabled(cfg))
+	})
+	t.Run("on", func(t *testing.T) {
+		cfg := new(config)
+		defaults(cfg, "", nil)
+		WithDBStats(1 * time.Second)
+		assert.True(t, int64(cfg.dbStats) == 1)
+		assert.True(t, dbStatsEnabled(cfg))
+	})
+	t.Run("interval 0", func(t *testing.T) {
+		// this test demonstrates that the logic for checking whether DBStats is enabled, is for the interval to be > 0
+		cfg := new(config)
+		defaults(cfg, "", nil)
+		WithDBStats(0 * time.Second)
+		assert.True(t, int64(cfg.dbStats) == 0)
+		assert.False(t, dbStatsEnabled(cfg))
+	})
+}

contrib/database/sql/sql.go

@@ -23,8 +23,10 @@ import (
 	"sync"
 	"time"
 
-	"gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql/internal"
+	sqlinternal "gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql/internal"
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"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/telemetry"
 )
@@ -154,7 +156,7 @@ func (t *tracedConnector) Connect(ctx context.Context) (driver.Conn, error) {
 		cfg:        t.cfg,
 	}
 	if dsn != "" {
-		tp.meta, _ = internal.ParseDSN(t.driverName, dsn)
+		tp.meta, _ = sqlinternal.ParseDSN(t.driverName, dsn)
 	}
 	start := time.Now()
 	ctx, end := startTraceTask(ctx, string(QueryTypeConnect))
@@ -209,7 +211,15 @@ func OpenDB(c driver.Connector, opts ...Option) *sql.DB {
 		driverName: driverName,
 		cfg:        cfg,
 	}
-	return sql.OpenDB(tc)
+	db := sql.OpenDB(tc)
+	if dbStatsEnabled(cfg) {
+		go pollDBStats(cfg.dbStats, db, pushDBStats)
+	}
+	return db
+}
+
+func dbStatsEnabled(cfg *config) bool {
+	return int64(cfg.dbStats) > 0
 }
 
 // Open returns connection to a DB using the traced version of the given driver. The driver may
@@ -250,3 +260,23 @@ func processOptions(cfg *config, driverName string, driver driver.Driver, dsn st
 	}
 	cfg.checkDBMPropagation(driverName, driver, dsn)
 }
+
+// pollDBStats calls (*DB).Stats on the db, at the specified interval. It pushes the DBStat off to the pushFn.
+func pollDBStats(interval time.Duration, db *sql.DB, pushFn func(stat sql.DBStats)) {
+	for range time.NewTicker(interval).C {
+		if db == nil {
+			log.Debug("No traced DB connection found; cannot pull DB stats.")
+			return
+		}
+		log.Debug("Traced DB connection found: DB stats will be gathered and sent every %v.", interval)
+		pushFn(db.Stats())
+	}
+}
+
+// pushDBStats separates the DBStats type out into individual statsd payloads and submits to the globalconfig's statsd client
+func pushDBStats(stats sql.DBStats) {
+	// Starting with just 1 metric & no tags, to complete a MVP.
+	openConns := stats.OpenConnections
+	s := internal.NewGauge("sql.db.open_connections", float64(openConns), nil, 1)
+	globalconfig.PushStat(s)
+}

contrib/database/sql/sql_test.go

@@ -7,6 +7,7 @@ package sql
 
 import (
 	"context"
+	"database/sql"
 	"database/sql/driver"
 	"errors"
 	"fmt"
@@ -44,6 +45,9 @@ func TestMain(m *testing.M) {
 	os.Exit(testResult)
 }
 
+// func TestDBStats(){}
+// Holding off on writing this until I get review on the implementation decision
+
 func TestSqlServer(t *testing.T) {
 	driverName := "sqlserver"
 	Register(driverName, &mssql.Driver{})
@@ -270,6 +274,32 @@ func TestOpenOptions(t *testing.T) {
 	})
 }
 
+func TestPollDBStats(t *testing.T) {
+	driverName := "postgres"
+	dsn := "postgres://postgres:[email protected]:5432/postgres?sslmode=disable"
+	db, err := Open(driverName, dsn)
+	assert.NoError(t, err)
+	interval := 3 * time.Millisecond
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		pollDBStats(interval, db, pollDBStatsCounter)
+	}()
+	time.Sleep(3 * interval)
+	assert.Len(t, dbStatsCollector, 3)
+}
+
+var dbStatsCollector []sql.DBStats
+
+func pollDBStatsCounter(stats sql.DBStats) {
+	dbStatsCollector = append(dbStatsCollector, stats)
+}
+
+// func TestPushDBStats(t *testing.T){
+
+// }
+
 func TestMySQLUint64(t *testing.T) {
 	Register("mysql", &mysql.MySQLDriver{})
 	defer unregister("mysql")

ddtrace/tracer/metrics_test.go

@@ -304,3 +304,24 @@ func TestTracerMetrics(t *testing.T) {
 	assert.Equal(1, calls["datadog.tracer.stopped"])
 	assert.True(tg.closed)
 }
+
+func TestContribStats(t *testing.T) {
+	// t.Run("gauge", func(t *testing.T) {
+	// 	assert := assert.New(t)
+	// 	var tg testStatsdClient
+	// 	trc := newUnstartedTracer(withStatsdClient(&tg))
+	// 	defer trc.statsd.Close()
+	// 	s := globalinternal.NewGauge("gauge", float64(1), nil, 1)
+	// 	tick := time.NewTicker(time.Millisecond)
+	// 	defer tick.Stop()
+	// 	trc.wg.Add(1)
+	// 	go func() {
+	// 		defer trc.wg.Done()
+	// 		for i := 0; i < 10; i++ {
+	// 			trc.statsCarrier.Add(s)
+	// 		}
+	// 		calls := tg.CallsByName()
+	// 		assert.Equal(10, calls["gauge"])
+	// 	}()
+	// })
+}

ddtrace/tracer/option.go

@@ -260,6 +260,8 @@ type config struct {
 
 	// headerAsTags holds the header as tags configuration.
 	headerAsTags dynamicConfig[[]string]
+
+	contribStats bool
 }
 
 // orchestrionConfig contains Orchestrion configuration.
@@ -340,6 +342,7 @@ func newConfig(opts ...StartOption) *config {
 		c.logToStdout = true
 	}
 	c.logStartup = internal.BoolEnv("DD_TRACE_STARTUP_LOGS", true)
+	c.contribStats = internal.BoolEnv("DD_TRACE_CONTRIB_STATS_ENABLED", true)
 	c.runtimeMetrics = internal.BoolEnv("DD_RUNTIME_METRICS_ENABLED", false)
 	c.debug = internal.BoolEnv("DD_TRACE_DEBUG", false)
 	c.enabled = internal.BoolEnv("DD_TRACE_ENABLED", true)
@@ -1258,6 +1261,15 @@ func setHeaderTags(headerAsTags []string) bool {
 	return true
 }
 
+// WithContribStats opens up a channel of communication between tracer and contrib libraries
+// for submitting stats from contribs to Datadog via the tracer's statsd client
+// It is enabled by default but can be disabled with `WithContribStats(false)`
+func WithContribStats(enabled bool) StartOption {
+	return func(c *config) {
+		c.contribStats = enabled
+	}
+}
+
 // UserMonitoringConfig is used to configure what is used to identify a user.
 // This configuration can be set by combining one or several UserMonitoringOption with a call to SetUser().
 type UserMonitoringConfig struct {

ddtrace/tracer/option_test.go

@@ -1262,6 +1262,27 @@ func TestWithHeaderTags(t *testing.T) {
 	assert.Equal(t, 0, globalconfig.HeaderTagsLen())
 }
 
+func TestContribStatsEnabled(t *testing.T) {
+	t.Run("Default", func(t *testing.T) {
+		c := newConfig()
+		assert.True(t, c.contribStats)
+	})
+	t.Run("Disable", func(t *testing.T) {
+		c := newConfig(WithContribStats(false))
+		assert.False(t, c.contribStats)
+	})
+	t.Run("Disable with ENV", func(t *testing.T) {
+		t.Setenv("DD_TRACE_CONTRIB_STATS_ENABLED", "false")
+		c := newConfig()
+		assert.False(t, c.contribStats)
+	})
+	t.Run("Env override", func(t *testing.T) {
+		t.Setenv("DD_TRACE_CONTRIB_STATS_ENABLED", "false")
+		c := newConfig(WithContribStats(true))
+		assert.True(t, c.contribStats)
+	})
+}
+
 func TestHostnameDisabled(t *testing.T) {
 	t.Run("DisabledWithUDS", func(t *testing.T) {
 		t.Setenv("DD_TRACE_AGENT_URL", "unix://somefakesocket")

ddtrace/tracer/tracer.go

@@ -23,6 +23,7 @@ import (
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/appsec"
 	appsecConfig "gopkg.in/DataDog/dd-trace-go.v1/internal/appsec/config"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/datastreams"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/hostname"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
 	"gopkg.in/DataDog/dd-trace-go.v1/internal/remoteconfig"
@@ -105,6 +106,8 @@ type tracer struct {
 	// abandonedSpansDebugger specifies where and how potentially abandoned spans are stored
 	// when abandoned spans debugging is enabled.
 	abandonedSpansDebugger *abandonedSpansDebugger
+
+	statsCarrier *globalinternal.StatsCarrier
 }
 
 const (
@@ -259,6 +262,10 @@ func newUnstartedTracer(opts ...StartOption) *tracer {
 			return f.DataStreams
 		})
 	}
+	var statsCarrier *globalinternal.StatsCarrier
+	if c.contribStats {
+		statsCarrier = globalinternal.NewStatsCarrier(statsd)
+	}
 	t := &tracer{
 		config:           c,
 		traceWriter:      writer,
@@ -278,8 +285,9 @@ func newUnstartedTracer(opts ...StartOption) *tracer {
 				Cache:            c.agent.HasFlag("sql_cache"),
 			},
 		}),
-		statsd:      statsd,
-		dataStreams: dataStreamsProcessor,
+		statsd:       statsd,
+		dataStreams:  dataStreamsProcessor,
+		statsCarrier: statsCarrier,
 	}
 	return t
 }
@@ -323,6 +331,10 @@ func newTracer(opts ...StartOption) *tracer {
 		t.reportHealthMetrics(statsInterval)
 	}()
 	t.stats.Start()
+	if sc := t.statsCarrier; sc != nil {
+		sc.Start()
+		globalconfig.SetStatsCarrier(sc)
+	}
 	return t
 }
 
@@ -651,6 +663,10 @@ func (t *tracer) Stop() {
 	if t.dataStreams != nil {
 		t.dataStreams.Stop()
 	}
+	if t.statsCarrier != nil {
+		t.statsCarrier.Stop()
+		globalconfig.ClearStatsCarrier()
+	}
 	appsec.Stop()
 	remoteconfig.Stop()
 }

ddtrace/tracer/tracer_test.go

@@ -31,6 +31,7 @@ import (
 	"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/telemetry"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/version"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -686,6 +687,52 @@ func TestTracerRuntimeMetrics(t *testing.T) {
 	})
 }
 
+var debugPrefix = fmt.Sprintf("Datadog Tracer %v DEBUG: ", version.Tag)
+
+func TestTracerContribStats(t *testing.T) {
+	stat := maininternal.NewGauge("test",float64(1),nil,1)
+	t.Run("default on", func(t *testing.T) {
+		tp := new(log.RecordLogger)
+		tracer := newTracer(WithDebugMode(true), WithLogger(tp))
+		defer tracer.Stop()
+		assert.NotNil(t, tracer.statsCarrier)
+		//check that the statscarrier has been set on the globalconfig
+		globalconfig.PushStat(stat)
+		assert.NotContains(t, tp.Logs(), fmt.Sprintf("%vNo stats carrier found; dropping stat %v", debugPrefix, stat.Name()))
+	})
+	t.Run("off", func(t *testing.T) {
+		tp := new(log.RecordLogger)
+		tracer := newTracer(WithContribStats(false), WithLogger(tp), WithDebugMode(true))
+		defer tracer.Stop()
+		assert.Nil(t, tracer.statsCarrier)
+		//check that there is no statscarrier on the globalconfig
+		globalconfig.PushStat(stat)
+		assert.Contains(t, tp.Logs(), fmt.Sprintf("%vNo stats carrier found; dropping stat %v", debugPrefix, stat.Name()))
+	})
+	t.Run("env", func(t *testing.T) {
+		os.Setenv("DD_TRACE_CONTRIB_STATS_ENABLED", "false")
+		defer os.Unsetenv("DD_TRACE_CONTRIB_STATS_ENABLED")
+		tp := new(log.RecordLogger)
+		tracer := newTracer(WithLogger(tp), WithDebugMode(true))
+		defer tracer.Stop()
+		assert.Nil(t, tracer.statsCarrier)
+		//check that there is no statscarrier on the globalconfig
+		globalconfig.PushStat(stat)
+		assert.Contains(t, tp.Logs(), fmt.Sprintf("%vNo stats carrier found; dropping stat %v", debugPrefix, stat.Name()))
+	})
+	t.Run("env override", func(t *testing.T) {
+		os.Setenv("DD_TRACE_CONTRIB_STATS_ENABLED", "false")
+		defer os.Unsetenv("DD_TRACE_CONTRIB_STATS_ENABLED")
+		tp := new(log.RecordLogger)
+		tracer := newTracer(WithLogger(tp), WithDebugMode(true), WithContribStats(true))
+		defer tracer.Stop()
+		assert.NotNil(t, tracer.statsCarrier)
+		//check that the statscarrier has been set on the globalconfig
+		globalconfig.PushStat(stat)
+		assert.NotContains(t, tp.Logs(), fmt.Sprintf("%vNo stats carrier found; dropping stat %v", debugPrefix, stat.Name()))
+	})
+}
+
 func TestTracerStartSpanOptions(t *testing.T) {
 	tracer := newTracer()
 	defer tracer.Stop()