fix: node_status exclusion check in SourceReaper.Run() for special metrics

Move the `isRoleExcluded` check earlier in the metric fetch loop to apply
consistently to all metric types, including special metrics like
`instance_up` and `change_events`.

Previously, special metrics weren't excluded if they have
a required `node_status` different than the server's

Author: 0xgouda

internal/reaper/source_reaper.go

@@ -145,6 +145,12 @@ func (sr *SourceReaper) Run(ctx context.Context) {
 			if lf := sr.lastFetch[name]; !lf.IsZero() && now.Sub(lf) < interval {
 				continue
 			}
+
+			metric, ok := metricDefs.GetMetricDef(name)
+			if !ok || sr.isRoleExcluded(metric) {
+				continue
+			}
+
 			switch {
 			case name == specialMetricServerLogEventCounts:
 				if sr.lastFetch[name].IsZero() {
@@ -157,16 +163,8 @@ func (sr *SourceReaper) Run(ctx context.Context) {
 			case IsDirectlyFetchableMetric(sr.md, name):
 				err = sr.fetchOSMetric(ctx, name)
 			case name == specialMetricChangeEvents || name == specialMetricInstanceUp:
-				err = sr.fetchSpecialMetric(ctx, name)
+				err = sr.fetchSpecialMetric(ctx, name, metric.StorageName)
 			default:
-				metric, ok := metricDefs.GetMetricDef(name)
-				if !ok {
-					l.WithField("metric", name).Warning("metric definition not found")
-					continue
-				}
-				if sr.isRoleExcluded(metric) {
-					continue
-				}
 				if cached := sr.reaper.GetMeasurementCache(sr.cacheKey(metric, name)); len(cached) > 0 {
 					l.WithField("metric", name).Info("instance level cache hit")
 					sr.sendEnvelope(ctx, name, metric.StorageName, cached)
@@ -264,14 +262,7 @@ func (sr *SourceReaper) fetchOSMetric(ctx context.Context, name string) error {
 }
 
 // fetchSpecialMetric handles change_events and instance_up metrics.
-func (sr *SourceReaper) fetchSpecialMetric(ctx context.Context, name string) error {
-	metric, ok := metricDefs.GetMetricDef(name)
-	if !ok {
-		return fmt.Errorf("metric definition not found for %s", name)
-	}
-	if sr.isRoleExcluded(metric) {
-		return nil
-	}
+func (sr *SourceReaper) fetchSpecialMetric(ctx context.Context, name, storageName string) error {
 	var (
 		data metrics.Measurements
 		err  error
@@ -286,7 +277,7 @@ func (sr *SourceReaper) fetchSpecialMetric(ctx context.Context, name string) err
 		return fmt.Errorf("failed to fetch special metric: %v", err)
 	}
 	if len(data) > 0 {
-		sr.sendEnvelope(ctx, name, metric.StorageName, data)
+		sr.sendEnvelope(ctx, name, storageName, data)
 	}
 	return err
 }

internal/reaper/source_reaper_integration_test.go

@@ -180,3 +180,103 @@ func TestIntegration_SourceReaper_RunCollectsMetrics(t *testing.T) {
 	assert.Equal(t, "integration_test", sMsg.DBName)
 	assert.NotEmpty(t, sMsg.Data)
 }
+
+func TestIntegration_SourceReaper_RunExcludesMetricsByNodeStatus(t *testing.T) {
+	md, tearDown := setupIntegrationDB(t)
+	defer tearDown()
+
+	helperSetNodeStatus := func(status string) {
+		metricDefs.MetricDefs["test_metric"] = metrics.Metric{
+			SQLs: metrics.SQLs{0: "SELECT 1 AS value"},
+			NodeStatus: status,
+		}
+		metricDefs.MetricDefs["server_log_event_counts"] = metrics.Metric{
+			SQLs: metrics.SQLs{0: "SELECT 1 AS value"},
+			NodeStatus: status,
+		}
+		metricDefs.MetricDefs["psutil_cpu"] = metrics.Metric{
+			SQLs: metrics.SQLs{0: "SELECT 1 AS value"},
+			NodeStatus: status,
+		}
+		metricDefs.MetricDefs[specialMetricInstanceUp] = metrics.Metric{
+			SQLs: metrics.SQLs{0: "SELECT 1 AS value"},
+			NodeStatus: status,
+		}
+	}
+	
+	r := &Reaper{
+		Options: &cmdopts.Options{
+			Metrics: metrics.CmdOpts{},
+			Sinks:   sinks.CmdOpts{},
+		},
+		measurementCh:    make(chan metrics.MeasurementEnvelope, 10),
+		measurementCache: NewInstanceMetricCache(),
+	}
+
+	// using psutil_*, server_log_event_counts, instance_up
+	// to ensure specially-handled metrics have the same behaviour
+	md.Metrics = metrics.MetricIntervals{
+		"test_metric": 5,
+		"server_log_event_counts": 5,
+		"psutil_cpu": 5,
+		specialMetricInstanceUp: 5,
+	}
+
+	t.Run("primary-only/standby-only metrics get excluded when node is standby/primary", func(t *testing.T) {
+		states := []string{"primary", "standby"}
+		for _, state := range states {
+			ctx, cancel := context.WithCancel(log.WithLogger(context.Background(), log.NewNoopLogger()))
+
+			md.IsInRecovery = true
+			if state == "standby" {
+				md.IsInRecovery = false
+			}
+
+			helperSetNodeStatus(state)
+
+			sr := NewSourceReaper(r, md)
+			go func() {
+				sr.Run(ctx)
+			}()
+
+			select {
+			case msg := <-r.measurementCh:
+				t.Errorf("Expected no measurement for primary-only metrics on standby, but got: %s", msg.MetricName)
+			case <-time.After(2 * time.Second):
+			}
+
+			cancel()
+		}
+	})
+
+	t.Run("primary-only/standby-only metrics get executed when node is primary/standby", func(t *testing.T) {
+		states := []string{"primary", "standby", ""} // "" => should fetch all as well
+		for _, state := range states {
+			ctx, cancel := context.WithCancel(log.WithLogger(context.Background(), log.NewNoopLogger()))
+
+			md.IsInRecovery = false
+			if state == "standby" {
+				md.IsInRecovery = true
+			}
+
+			helperSetNodeStatus(state)
+
+			sr := NewSourceReaper(r, md)
+			go func() {
+				sr.Run(ctx)
+			}()
+
+			time.Sleep(2 * time.Second)
+			assert.GreaterOrEqual(t, len(r.measurementCh), 3)
+			cancel()
+
+			for range len(r.measurementCh) {
+				// empty channel to ensure correctness in subsequent runs
+				select {
+				case <-r.measurementCh:
+				default:
+				}
+			}
+		}
+	})
+}

internal/reaper/source_reaper_test.go

@@ -307,52 +307,12 @@ func TestSourceReaper_FetchSpecialMetric(t *testing.T) {
 		return NewSourceReaper(r, md), md, mock
 	}
 
-	t.Run("metric not found in definitions", func(t *testing.T) {
-		sr, _, mock := newSR(t)
-		defer mock.Close()
-		assert.Error(t, sr.fetchSpecialMetric(ctx, "no_such_special_xyz"))
-		select {
-		case <-sr.reaper.measurementCh:
-			t.Error("expected no measurement")
-		default:
-		}
-		assert.NoError(t, mock.ExpectationsWereMet())
-	})
-
-	t.Run("primary-only metric skipped on standby", func(t *testing.T) {
-		sr, md, mock := newSR(t)
-		defer mock.Close()
-		metricDefs.MetricDefs["sp_primary_only"] = metrics.Metric{NodeStatus: "primary"}
-		md.IsInRecovery = true
-		assert.NoError(t, sr.fetchSpecialMetric(ctx, "sp_primary_only"))
-		select {
-		case <-sr.reaper.measurementCh:
-			t.Error("expected no measurement for primary-only on standby")
-		default:
-		}
-		assert.NoError(t, mock.ExpectationsWereMet())
-	})
-
-	t.Run("standby-only metric skipped on primary", func(t *testing.T) {
-		sr, md, mock := newSR(t)
-		defer mock.Close()
-		metricDefs.MetricDefs["sp_standby_only"] = metrics.Metric{NodeStatus: "standby"}
-		md.IsInRecovery = false
-		assert.NoError(t, sr.fetchSpecialMetric(ctx, "sp_standby_only"))
-		select {
-		case <-sr.reaper.measurementCh:
-			t.Error("expected no measurement for standby-only on primary")
-		default:
-		}
-		assert.NoError(t, mock.ExpectationsWereMet())
-	})
+	sr, _, mock := newSR(t)
+	defer mock.Close()
 
 	t.Run("instance_up dispatches measurement on ping success", func(t *testing.T) {
-		sr, _, mock := newSR(t)
-		defer mock.Close()
-		metricDefs.MetricDefs[specialMetricInstanceUp] = metrics.Metric{}
 		mock.ExpectPing()
-		assert.NoError(t, sr.fetchSpecialMetric(ctx, specialMetricInstanceUp))
+		assert.NoError(t, sr.fetchSpecialMetric(ctx, specialMetricInstanceUp, ""))
 		select {
 		case msg := <-sr.reaper.measurementCh:
 			assert.Equal(t, specialMetricInstanceUp, msg.MetricName)
@@ -365,11 +325,8 @@ func TestSourceReaper_FetchSpecialMetric(t *testing.T) {
 	})
 
 	t.Run("instance_up uses storage name when set", func(t *testing.T) {
-		sr, _, mock := newSR(t)
-		defer mock.Close()
-		metricDefs.MetricDefs[specialMetricInstanceUp] = metrics.Metric{StorageName: "infra_up"}
 		mock.ExpectPing()
-		assert.NoError(t, sr.fetchSpecialMetric(ctx, specialMetricInstanceUp))
+		assert.NoError(t, sr.fetchSpecialMetric(ctx, specialMetricInstanceUp, "infra_up"))
 		select {
 		case msg := <-sr.reaper.measurementCh:
 			assert.Equal(t, "infra_up", msg.MetricName)
@@ -380,13 +337,9 @@ func TestSourceReaper_FetchSpecialMetric(t *testing.T) {
 	})
 
 	t.Run("change_events dispatches no measurement when no hash defs present", func(t *testing.T) {
-		sr, _, mock := newSR(t)
-		defer mock.Close()
+		// Doesn't contain additional defs for any of {"sproc_hashes", "table_hashes", "index_hashes", "configuration_hashes", "privilege_hashes"}
 		metricDefs.MetricDefs[specialMetricChangeEvents] = metrics.Metric{}
-		for _, name := range []string{"sproc_hashes", "table_hashes", "index_hashes", "configuration_hashes", "privilege_hashes"} {
-			delete(metricDefs.MetricDefs, name)
-		}
-		assert.NoError(t, sr.fetchSpecialMetric(ctx, specialMetricChangeEvents))
+		assert.NoError(t, sr.fetchSpecialMetric(ctx, specialMetricChangeEvents, ""))
 		select {
 		case <-sr.reaper.measurementCh:
 			t.Error("expected no measurement when no changes detected")