Logic for metric override mode

PiperOrigin-RevId: 723078387
GoogleCloudPlatform · Feb 4, 2025 · 52473ed · 52473ed
1 parent 254cfc0
commit 52473ed
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 15 deletions.
diff --git a/internal/daemon/daemon.go b/internal/daemon/daemon.go
@@ -153,6 +153,28 @@ func (d *Daemon) startdaemonHandler(ctx context.Context, cancel context.CancelFu
 		return err
 	}
 
+	// Check if the metric override file exists. If it does, operate in override mode.
+	// Override mode will collect metrics from the override file and send them to Data Warehouse.
+	// Override mode will not start any other services.
+	if fileInfo, err := os.ReadFile(workloadmanager.MetricOverridePath); fileInfo != nil && err == nil {
+		log.Logger.Info("Metric override file found. Operating in override mode.")
+		metricCollectionService := workloadmanager.Service{Config: d.config, Client: wlmClient}
+		metricCollectionCtx := log.SetCtx(ctx, "context", "WorkloadManagerMetrics")
+		recoverableStart := &recovery.RecoverableRoutine{
+			Routine:             metricCollectionService.CollectAndSendMetricsToDataWarehouse,
+			RoutineArg:          d.config,
+			ErrorCode:           0,
+			ExpectedMinDuration: 0,
+			UsageLogger:         *usagemetrics.UsageLogger,
+		}
+		recoverableStart.StartRoutine(metricCollectionCtx)
+
+		// Log a RUNNING usage metric once a day.
+		go usagemetrics.LogRunningDaily()
+		d.waitForShutdown(shutdownch, cancel)
+		return nil
+	}
+
 	log.Logger.Info("Starting common discovery")
 	oracleCh := make(chan *servicecommunication.Message, 3)
 	mySQLCh := make(chan *servicecommunication.Message, 3)
@@ -202,17 +224,6 @@ func (d *Daemon) startdaemonHandler(ctx context.Context, cancel context.CancelFu
 		recoverableStart.StartRoutine(ctx)
 	}
 
-	metricCollectionService := workloadmanager.Service{Config: d.config, Client: wlmClient}
-	metricCollectionCtx := log.SetCtx(ctx, "context", "WorkloadManagerMetrics")
-	recoverableStart = &recovery.RecoverableRoutine{
-		Routine:             metricCollectionService.CollectAndSendMetricsToDataWarehouse,
-		RoutineArg:          d.config,
-		ErrorCode:           0,
-		ExpectedMinDuration: 0,
-		UsageLogger:         *usagemetrics.UsageLogger,
-	}
-	recoverableStart.StartRoutine(metricCollectionCtx)
-
 	// Log a RUNNING usage metric once a day.
 	go usagemetrics.LogRunningDaily()
 	d.waitForShutdown(shutdownch, cancel)

diff --git a/internal/workloadmanager/workloadcollector.go b/internal/workloadmanager/workloadcollector.go
@@ -90,7 +90,8 @@ type Service struct {
 	Client WLMWriter
 }
 
-const metricOverridePath = "/etc/google-cloud-workload-agent/wlmmetricoverride.yaml"
+// MetricOverridePath is the path to the metric override file.
+const MetricOverridePath = "/etc/google-cloud-workload-agent/wlmmetricoverride.yaml"
 
 // Client creates a new WLM client.
 func Client(ctx context.Context, config *cpb.Configuration) (WLMWriter, error) {
@@ -115,7 +116,7 @@ func (s *Service) CollectAndSendMetricsToDataWarehouse(ctx context.Context, a an
 }
 
 func readAndLogMetricOverrideYAML(ctx context.Context, reader ConfigFileReader) bool {
-	file, err := reader(metricOverridePath)
+	file, err := reader(MetricOverridePath)
 	if err != nil {
 		log.CtxLogger(ctx).Debugw("Could not read the metric override file", "error", err)
 		return false
@@ -138,7 +139,7 @@ func readAndLogMetricOverrideYAML(ctx context.Context, reader ConfigFileReader)
 
 // collectOverrideMetrics reads workload metrics from an override file.
 func collectOverrideMetrics(ctx context.Context, reader ConfigFileReader) []WorkloadMetrics {
-	file, err := reader(metricOverridePath)
+	file, err := reader(MetricOverridePath)
 	if err != nil {
 		log.CtxLogger(ctx).Debugw("Could not read the metric override file", "error", err)
 		return []WorkloadMetrics{}
@@ -263,7 +264,7 @@ func createWriteInsightRequest(ctx context.Context, wm WorkloadMetrics, cp *cpb.
 				ValidationDetails: wm.Metrics,
 				ProjectId:         cp.GetProjectId(),
 				InstanceName:      cp.GetInstanceName(),
-				AgentVersion: configuration.AgentVersion,
+				AgentVersion:      configuration.AgentVersion,
 			},
 		},
 	}