fix(database_observability.postgres): cleanup embedded exporter collectors on reconnection

This is a port of #6021

Author: cristiangreco

internal/component/database_observability/postgres/component.go

@@ -291,6 +291,8 @@ func (c *Component) Run(ctx context.Context) error {
 			for _, collector := range c.collectors {
 				collector.Stop()
 			}
+			c.cleanupExporterCollectors()
+
 			if c.dbConnection != nil {
 				c.dbConnection.Close()
 			}
@@ -367,11 +369,24 @@ func (c *Component) tryReconnect(ctx context.Context) error {
 	return nil
 }
 
+// cleanupExporterCollectors releases resources held by embedded exporter collectors.
+// Callers must hold c.mut.
+func (c *Component) cleanupExporterCollectors() {
+	for _, col := range c.exporterCollectors {
+		if closable, ok := col.(interface{ CloseServers() }); ok {
+			closable.CloseServers()
+		}
+		c.registry.Unregister(col)
+	}
+	c.exporterCollectors = nil
+}
+
 func (c *Component) connectAndStartCollectors(ctx context.Context) error {
 	if c.dbConnection != nil {
 		c.dbConnection.Close()
 		c.dbConnection = nil
 	}
+	c.cleanupExporterCollectors()
 
 	dbConnection, err := c.openSQL("postgres", string(c.args.DataSourceName))
 	if err != nil {
@@ -415,11 +430,6 @@ func (c *Component) connectAndStartCollectors(ctx context.Context) error {
 		cp = cloudProvider
 	}
 
-	for _, col := range c.exporterCollectors {
-		c.registry.Unregister(col)
-	}
-	c.exporterCollectors = nil
-
 	if len(c.args.Targets) == 0 {
 		if c.args.PrometheusExporter == nil {
 			d := PrometheusExporterArguments(exporter_postgres.DefaultArguments)

internal/component/database_observability/postgres/component_test.go

@@ -26,6 +26,43 @@ import (
 	"github.com/grafana/loki/pkg/push"
 )
 
+func newTestComponent(t *testing.T, openSQL func(string, string) (*sql.DB, error)) *Component {
+	t.Helper()
+	opts := cmp.Options{
+		ID:            "test",
+		Logger:        kitlog.NewNopLogger(),
+		OnStateChange: func(e cmp.Exports) {},
+		GetServiceData: func(name string) (any, error) {
+			return http_service.Data{MemoryListenAddr: "127.0.0.1:0", BaseHTTPPath: "/"}, nil
+		},
+	}
+	args := Arguments{
+		DataSourceName:    alloytypes.Secret("postgres://user:pass@127.0.0.1:5432/db?sslmode=disable"),
+		ForwardTo:         []loki.LogsReceiver{},
+		Targets:           []discovery.Target{},
+		DisableCollectors: []string{"query_details", "schema_details", "query_samples", "explain_plans"},
+		HealthCheckArguments: HealthCheckArguments{
+			CollectInterval: 1 * time.Hour,
+		},
+	}
+	c := &Component{
+		opts:         opts,
+		args:         args,
+		fanout:       loki.NewFanout(args.ForwardTo),
+		handler:      loki.NewLogsReceiver(),
+		registry:     prometheus.NewRegistry(),
+		healthErr:    atomic.NewString(""),
+		openSQL:      openSQL,
+		logsReceiver: loki.NewLogsReceiver(),
+	}
+	c.instanceKey = "test-instance"
+	c.baseTarget = discovery.NewTargetFromMap(map[string]string{
+		"instance": c.instanceKey,
+		"job":      "database_observability",
+	})
+	return c
+}
+
 func Test_enableOrDisableCollectors(t *testing.T) {
 	t.Run("nothing specified (default behavior)", func(t *testing.T) {
 		exampleDBO11yAlloyConfig := `
@@ -643,6 +680,62 @@ func Test_connectAndStartCollectors(t *testing.T) {
 	})
 }
 
+type fakeClosableCollector struct {
+	prometheus.Collector
+	closeCalls int
+}
+
+func newFakeClosableCollector(name string) *fakeClosableCollector {
+	gauge := prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: name,
+		Help: name,
+	})
+	return &fakeClosableCollector{Collector: gauge}
+}
+
+func (c *fakeClosableCollector) CloseServers() {
+	c.closeCalls++
+}
+
+func TestComponent_cleanupExporterCollectors(t *testing.T) {
+	t.Run("closes closable exporters and unregisters them", func(t *testing.T) {
+		registry := prometheus.NewRegistry()
+		collector := newFakeClosableCollector("test_cleanup_exporter_collectors_closable")
+		require.NoError(t, registry.Register(collector))
+
+		c := &Component{
+			registry:           registry,
+			exporterCollectors: []prometheus.Collector{collector},
+		}
+
+		c.cleanupExporterCollectors()
+
+		assert.Equal(t, 1, collector.closeCalls)
+		assert.Nil(t, c.exporterCollectors)
+		assert.False(t, registry.Unregister(collector))
+	})
+
+	t.Run("unregisters non-closable collectors without panicking", func(t *testing.T) {
+		registry := prometheus.NewRegistry()
+		collector := prometheus.NewGauge(prometheus.GaugeOpts{
+			Name: "test_cleanup_exporter_collectors_plain",
+			Help: "test",
+		})
+		collector.Set(1)
+		require.NoError(t, registry.Register(collector))
+
+		c := &Component{
+			registry:           registry,
+			exporterCollectors: []prometheus.Collector{collector},
+		}
+
+		c.cleanupExporterCollectors()
+
+		assert.Nil(t, c.exporterCollectors)
+		assert.False(t, registry.Unregister(collector))
+	})
+}
+
 func TestPostgres_Reconnection(t *testing.T) {
 	t.Run("tryReconnect fails and maintains health error", func(t *testing.T) {
 		opts := cmp.Options{
@@ -671,47 +764,14 @@ func TestPostgres_Reconnection(t *testing.T) {
 	})
 
 	t.Run("tryReconnect succeeds and clears health error", func(t *testing.T) {
-		opts := cmp.Options{
-			ID:            "test",
-			Logger:        kitlog.NewNopLogger(),
-			OnStateChange: func(e cmp.Exports) {},
-			GetServiceData: func(name string) (any, error) {
-				return http_service.Data{MemoryListenAddr: "127.0.0.1:0", BaseHTTPPath: "/"}, nil
-			},
-		}
-
-		args := Arguments{
-			DataSourceName:    alloytypes.Secret("postgres://user:pass@127.0.0.1:5432/db?sslmode=disable"),
-			ForwardTo:         []loki.LogsReceiver{},
-			Targets:           []discovery.Target{},
-			DisableCollectors: []string{"query_details", "schema_details", "query_samples", "explain_plans"},
-			HealthCheckArguments: HealthCheckArguments{
-				CollectInterval: 1 * time.Hour,
-			},
-		}
-
 		// First mock: will fail
 		db1, mock1, err := sqlmock.New(sqlmock.MonitorPingsOption(true))
 		require.NoError(t, err)
 		defer db1.Close()
 
 		mock1.ExpectPing().WillReturnError(assert.AnError)
 
-		c := &Component{
-			opts:         opts,
-			args:         args,
-			fanout:       loki.NewFanout(args.ForwardTo),
-			handler:      loki.NewLogsReceiver(),
-			registry:     prometheus.NewRegistry(),
-			healthErr:    atomic.NewString(""),
-			openSQL:      func(_ string, _ string) (*sql.DB, error) { return db1, nil },
-			logsReceiver: loki.NewLogsReceiver(),
-		}
-		c.instanceKey = "test-instance"
-		c.baseTarget = discovery.NewTargetFromMap(map[string]string{
-			"instance": c.instanceKey,
-			"job":      "database_observability",
-		})
+		c := newTestComponent(t, func(_, _ string) (*sql.DB, error) { return db1, nil })
 
 		// First attempt: connection fails
 		err = c.tryReconnect(context.Background())
@@ -722,12 +782,10 @@ func TestPostgres_Reconnection(t *testing.T) {
 		db2, mock2, err := sqlmock.New(sqlmock.MonitorPingsOption(true))
 		require.NoError(t, err)
 		defer db2.Close()
-
 		mock2.ExpectPing()
 		mock2.ExpectQuery(`SELECT.*system_identifier.*inet_server_addr.*inet_server_port.*version`).
 			WillReturnRows(sqlmock.NewRows([]string{"system_identifier", "inet_server_addr", "inet_server_port", "version"}).
 				AddRow("1234567890", "127.0.0.1", "5432", "14.0"))
-
 		c.openSQL = func(_ string, _ string) (*sql.DB, error) { return db2, nil }
 
 		// Second attempt: connection succeeds and clears error
@@ -760,21 +818,27 @@ func TestPostgres_Reconnection(t *testing.T) {
 		require.NoError(t, err)
 
 		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
 
-		runErr := make(chan error, 1)
-		go func() {
-			runErr <- c.Run(ctx)
-		}()
+		err = c.Run(ctx)
+		assert.NoError(t, err)
+	})
+
+	t.Run("Run cleans up embedded exporter collectors on shutdown", func(t *testing.T) {
+		c := newTestComponent(t, func(_, _ string) (*sql.DB, error) { return nil, assert.AnError })
+		oldCollector := newFakeClosableCollector("test_run_cleanup_old_exporter")
+		require.NoError(t, c.registry.Register(oldCollector))
+		c.exporterCollectors = []prometheus.Collector{oldCollector}
 
-		time.Sleep(100 * time.Millisecond)
+		ctx, cancel := context.WithCancel(context.Background())
 		cancel()
 
-		select {
-		case err := <-runErr:
-			assert.NoError(t, err)
-		case <-time.After(5 * time.Second):
-			t.Fatal("Run did not exit after context cancellation")
-		}
+		err := c.Run(ctx)
+		assert.NoError(t, err)
+
+		assert.Equal(t, 1, oldCollector.closeCalls)
+		assert.Nil(t, c.exporterCollectors)
+		assert.False(t, c.registry.Unregister(oldCollector))
 	})
 }