Revert "Separate baseline path scheduling"

This reverts commit c5ec567.

Author: AlexandreYang

comp/networkpath/npcollector/impl/baseline.go

@@ -6,12 +6,9 @@
 package npcollectorimpl
 
 import (
-	"iter"
 	"sort"
 
 	"github.com/DataDog/datadog-agent/comp/networkpath/npcollector/impl/common"
-	npmodel "github.com/DataDog/datadog-agent/comp/networkpath/npcollector/model"
-	"github.com/DataDog/datadog-agent/pkg/networkpath/payload"
 )
 
 const baselineSelectionsPerSnapshot = 3
@@ -59,36 +56,3 @@ func addBaselineCandidate(selected []baselineCandidate, candidate baselineCandid
 	sort.Slice(selected, func(i, j int) bool { return selected[i].betterThan(selected[j]) })
 	return selected
 }
-
-func (s *npCollectorImpl) scheduleBaselineNetworkPathTests(conns iter.Seq[npmodel.NetworkPathConnection]) {
-	vpcSubnets, err := s.getVPCSubnets()
-	if err != nil {
-		s.logger.Errorf("Failed to get VPC subnets to skip: %s", err)
-		return
-	}
-
-	startTime := s.TimeNowFn()
-	connCount := 0
-	selected := make([]baselineCandidate, 0, baselineSelectionsPerSnapshot)
-	for conn := range conns {
-		connCount++
-		pathtest, _, ok := s.prepareNetworkPathTest(conn, payload.PathOriginNetworkTraffic, vpcSubnets)
-		if !ok {
-			continue
-		}
-		pathtest.DynamicTestProfile = payload.DynamicTestProfileBaseline
-		selected = addBaselineCandidate(selected, baselineCandidate{
-			path:       pathtest,
-			pathHash:   pathtest.GetHash(),
-			diagnostic: conn.Baseline.Diagnostic,
-			bytes:      conn.Baseline.Bytes,
-		})
-	}
-
-	for i := range selected {
-		if err := s.scheduleOne(&selected[i].path); err != nil {
-			s.logger.Errorf("Error scheduling baseline pathtest: %s", err)
-		}
-	}
-	s.reportScheduleTelemetry(startTime, connCount)
-}

comp/networkpath/npcollector/impl/npcollector.go

@@ -286,35 +286,22 @@ func (s *npCollectorImpl) getVPCSubnets() ([]netip.Prefix, error) {
 }
 
 func (s *npCollectorImpl) ScheduleNetworkPathTests(conns iter.Seq[npmodel.NetworkPathConnection]) {
-	if s.collectorConfigs.connectionsMonitoringEnabled {
-		s.scheduleNetworkPathTests(payload.PathOriginNetworkTraffic, conns)
-	} else if s.collectorConfigs.baselineTestsEnabled {
-		s.scheduleBaselineNetworkPathTests(conns)
+	if !s.collectorConfigs.connectionsMonitoringEnabled && !s.collectorConfigs.baselineTestsEnabled {
+		return
 	}
+
+	baseline := !s.collectorConfigs.connectionsMonitoringEnabled
+	s.scheduleNetworkPathTests(payload.PathOriginNetworkTraffic, conns, baseline)
 }
 
 func (s *npCollectorImpl) ScheduleNetflowPathTests(conns iter.Seq[npmodel.NetworkPathConnection]) {
 	if !s.collectorConfigs.netflowMonitoringEnabled {
 		return
 	}
-	s.scheduleNetworkPathTests(payload.PathOriginNetflow, conns)
-}
-
-func (s *npCollectorImpl) prepareNetworkPathTest(conn npmodel.NetworkPathConnection, origin payload.PathOrigin, vpcSubnets []netip.Prefix) (common.Pathtest, pathEvaluation, bool) {
-	evaluation := s.evaluateNetworkPathForConn(conn, origin, vpcSubnets)
-	if !evaluation.shouldSchedule {
-		s.logger.Tracef("Skipped connection: addr=%s, protocol=%s", conn.Dest, conn.Type)
-		return common.Pathtest{}, pathEvaluation{}, false
-	}
-	return s.makePathtest(conn, origin), evaluation, true
-}
-
-func (s *npCollectorImpl) reportScheduleTelemetry(startTime time.Time, connCount int) {
-	_ = s.statsdClient.Count(common.NetworkPathCollectorMetricPrefix+"schedule.conns_received", int64(connCount), []string{}, 1)
-	_ = s.statsdClient.Gauge(common.NetworkPathCollectorMetricPrefix+"schedule.duration", s.TimeNowFn().Sub(startTime).Seconds(), nil, 1)
+	s.scheduleNetworkPathTests(payload.PathOriginNetflow, conns, false)
 }
 
-func (s *npCollectorImpl) scheduleNetworkPathTests(origin payload.PathOrigin, conns iter.Seq[npmodel.NetworkPathConnection]) {
+func (s *npCollectorImpl) scheduleNetworkPathTests(origin payload.PathOrigin, conns iter.Seq[npmodel.NetworkPathConnection], baseline bool) {
 	var vpcSubnets []netip.Prefix
 	if origin == payload.PathOriginNetworkTraffic {
 		var err error
@@ -327,22 +314,46 @@ func (s *npCollectorImpl) scheduleNetworkPathTests(origin payload.PathOrigin, co
 
 	startTime := s.TimeNowFn()
 	connCount := 0
+	var selected []baselineCandidate
+	if baseline {
+		selected = make([]baselineCandidate, 0, baselineSelectionsPerSnapshot)
+	}
 	for conn := range conns {
 		connCount++
-		pathtest, evaluation, ok := s.prepareNetworkPathTest(conn, origin, vpcSubnets)
-		if !ok {
+		evaluation := s.evaluateNetworkPathForConn(conn, origin, vpcSubnets)
+		if !evaluation.shouldSchedule {
+			s.logger.Tracef("Skipped connection: addr=%s, protocol=%s", conn.Dest, conn.Type)
+			continue
+		}
+		pathtest := s.makePathtest(conn, origin)
+		if baseline {
+			pathtest.DynamicTestProfile = payload.DynamicTestProfileBaseline
+			selected = addBaselineCandidate(selected, baselineCandidate{
+				path:       pathtest,
+				pathHash:   pathtest.GetHash(),
+				diagnostic: conn.Baseline.Diagnostic,
+				bytes:      conn.Baseline.Bytes,
+			})
 			continue
 		}
+
 		pathtest.TestConfigID = evaluation.testConfigID
 		pathtest.Tags = evaluation.tags
 		if evaluation.testConfigID != "" {
 			pathtest.TestConfigSource = payload.TestConfigSourceRemote
 		}
-		if err := s.scheduleOne(&pathtest); err != nil {
+		err := s.scheduleOne(&pathtest)
+		if err != nil {
 			s.logger.Errorf("Error scheduling pathtests: %s", err)
 		}
 	}
-	s.reportScheduleTelemetry(startTime, connCount)
+	for i := range selected {
+		if err := s.scheduleOne(&selected[i].path); err != nil {
+			s.logger.Errorf("Error scheduling baseline pathtest: %s", err)
+		}
+	}
+	_ = s.statsdClient.Count(common.NetworkPathCollectorMetricPrefix+"schedule.conns_received", int64(connCount), []string{}, 1)
+	_ = s.statsdClient.Gauge(common.NetworkPathCollectorMetricPrefix+"schedule.duration", s.TimeNowFn().Sub(startTime).Seconds(), nil, 1)
 }
 
 // scheduleOne schedules pathtests.