feat: monitor sql database (#444)

sqlutil/monitor.go

@@ -0,0 +1,54 @@
+package sqlutil
+
+import (
+	"context"
+	"database/sql"
+	"time"
+
+	"github.com/rudderlabs/rudder-go-kit/config"
+	"github.com/rudderlabs/rudder-go-kit/stats"
+)
+
+// MonitorDatabase collects database connection pool metrics at regular intervals synchronously until the context is canceled.
+func MonitorDatabase(
+	ctx context.Context,
+	conf *config.Config,
+	statsFactory stats.Stats,
+	db *sql.DB,
+	identifier string,
+) {
+	statsReportInterval := conf.GetDurationVar(10, time.Second, "Database.statsReportInterval")
+
+	tags := stats.Tags{
+		"identifier": identifier,
+	}
+
+	maxOpenConnectionsStat := statsFactory.NewTaggedStat("db_max_open_connections", stats.GaugeType, tags)
+	openConnectionsStat := statsFactory.NewTaggedStat("db_open_connections", stats.GaugeType, tags)
+	inUseStat := statsFactory.NewTaggedStat("db_in_use", stats.GaugeType, tags)
+	idleStat := statsFactory.NewTaggedStat("db_idle", stats.GaugeType, tags)
+	waitCountStat := statsFactory.NewTaggedStat("db_wait_count", stats.GaugeType, tags)
+	waitDurationStat := statsFactory.NewTaggedStat("db_wait_duration", stats.TimerType, tags)
+	maxIdleClosedStat := statsFactory.NewTaggedStat("db_max_idle_closed", stats.GaugeType, tags)
+	maxIdleTimeClosedStat := statsFactory.NewTaggedStat("db_max_idle_time_closed", stats.GaugeType, tags)
+	maxLifetimeClosedStat := statsFactory.NewTaggedStat("db_max_lifetime_closed", stats.GaugeType, tags)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(statsReportInterval):
+			dbStats := db.Stats()
+
+			maxOpenConnectionsStat.Gauge(dbStats.MaxOpenConnections)
+			openConnectionsStat.Gauge(dbStats.OpenConnections)
+			inUseStat.Gauge(dbStats.InUse)
+			idleStat.Gauge(dbStats.Idle)
+			waitCountStat.Gauge(int(dbStats.WaitCount))
+			waitDurationStat.SendTiming(dbStats.WaitDuration)
+			maxIdleClosedStat.Gauge(int(dbStats.MaxIdleClosed))
+			maxIdleTimeClosedStat.Gauge(int(dbStats.MaxIdleTimeClosed))
+			maxLifetimeClosedStat.Gauge(int(dbStats.MaxLifetimeClosed))
+		}
+	}
+}

sqlutil/monitor_test.go

@@ -0,0 +1,71 @@
+package sqlutil_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/ory/dockertest/v3"
+	"github.com/stretchr/testify/require"
+
+	"github.com/rudderlabs/rudder-go-kit/config"
+	"github.com/rudderlabs/rudder-go-kit/sqlutil"
+	"github.com/rudderlabs/rudder-go-kit/stats"
+	"github.com/rudderlabs/rudder-go-kit/stats/memstats"
+	"github.com/rudderlabs/rudder-go-kit/testhelper/docker/resource/postgres"
+)
+
+func TestMonitorDatabase(t *testing.T) {
+	pool, err := dockertest.NewPool("")
+	require.NoError(t, err)
+	postgresContainer, err := postgres.Setup(pool, t)
+	require.NoError(t, err)
+
+	postgresContainer.DB.SetMaxOpenConns(10)
+	postgresContainer.DB.SetMaxIdleConns(5)
+
+	statsStore, err := memstats.New()
+	require.NoError(t, err)
+
+	identifier := "test"
+
+	conf := config.New()
+	conf.Set("Database.statsReportInterval", "1s")
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	setupCh := make(chan struct{})
+	go func() {
+		defer close(setupCh)
+		sqlutil.MonitorDatabase(ctx, conf, statsStore, postgresContainer.DB, identifier)
+	}()
+
+	require.Eventually(t, func() bool {
+		return statsStore.Get("db_max_open_connections", stats.Tags{
+			"identifier": identifier,
+		}).LastValue() == 10
+	},
+		5*time.Second,
+		100*time.Millisecond,
+	)
+	require.Eventually(t, func() bool {
+		return statsStore.Get("db_open_connections", stats.Tags{
+			"identifier": identifier,
+		}).LastValue() == 1
+	},
+		5*time.Second,
+		100*time.Millisecond,
+	)
+	require.Eventually(t, func() bool {
+		return statsStore.Get("db_idle", stats.Tags{
+			"identifier": identifier,
+		}).LastValue() == 1
+	},
+		5*time.Second,
+		100*time.Millisecond,
+	)
+
+	cancel()
+	<-setupCh
+}