From 28c95d2129130134ee4c3e50cc36c1694d990263 Mon Sep 17 00:00:00 2001
From: Vihas Makwana <vihas.makwana@elastic.co>
Date: Thu, 27 Feb 2025 20:35:05 +0530
Subject: [PATCH] add options for registry

---
 report/setup.go | 42 ++++++++++++++++++++++++++++++++++--------
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/report/setup.go b/report/setup.go
index 3406a0dc56..43cfca4aae 100644
--- a/report/setup.go
+++ b/report/setup.go
@@ -40,14 +40,40 @@ func init() {
 	systemMetrics = monitoring.Default.NewRegistry("system")
 }
 
+type option struct {
+	systemMetrics  *monitoring.Registry
+	processMetrics *monitoring.Registry
+}
+
+type OptionFunc func(o *option)
+
+func WithProcessRegistry(r *monitoring.Registry) OptionFunc {
+	return func(o *option) {
+		o.processMetrics = r
+	}
+}
+
+func WithSystemRegistry(r *monitoring.Registry) OptionFunc {
+	return func(o *option) {
+		o.systemMetrics = r
+	}
+}
+
 // monitoringCgroupsHierarchyOverride is an undocumented environment variable which
 // overrides the cgroups path under /sys/fs/cgroup, which should be set to "/" when running
 // Elastic Agent under Docker.
 const monitoringCgroupsHierarchyOverride = "LIBBEAT_MONITORING_CGROUPS_HIERARCHY_OVERRIDE"
 
 // SetupMetrics creates a basic suite of metrics handlers for monitoring, including build info and system resources
-func SetupMetrics(logger *logp.Logger, name, version string) error {
-	monitoring.NewFunc(systemMetrics, "cpu", ReportSystemCPUUsage, monitoring.Report)
+func SetupMetrics(logger *logp.Logger, name, version string, opts ...OptionFunc) error {
+	opt := &option{
+		systemMetrics:  systemMetrics,
+		processMetrics: processMetrics,
+	}
+	for _, o := range opts {
+		o(opt)
+	}
+	monitoring.NewFunc(opt.systemMetrics, "cpu", ReportSystemCPUUsage, monitoring.Report)
 
 	name = processName(name)
 	processStats = &process.Stats{
@@ -63,12 +89,12 @@ func SetupMetrics(logger *logp.Logger, name, version string) error {
 		return fmt.Errorf("failed to init process stats for agent: %w", err)
 	}
 
-	monitoring.NewFunc(processMetrics, "memstats", MemStatsReporter(logger, processStats), monitoring.Report)
-	monitoring.NewFunc(processMetrics, "cpu", InstanceCPUReporter(logger, processStats), monitoring.Report)
-	monitoring.NewFunc(processMetrics, "runtime", ReportRuntime, monitoring.Report)
-	monitoring.NewFunc(processMetrics, "info", infoReporter(name, version), monitoring.Report)
+	monitoring.NewFunc(opt.processMetrics, "memstats", MemStatsReporter(logger, processStats), monitoring.Report)
+	monitoring.NewFunc(opt.processMetrics, "cpu", InstanceCPUReporter(logger, processStats), monitoring.Report)
+	monitoring.NewFunc(opt.processMetrics, "runtime", ReportRuntime, monitoring.Report)
+	monitoring.NewFunc(opt.processMetrics, "info", infoReporter(name, version), monitoring.Report)
 
-	setupPlatformSpecificMetrics(logger, processStats)
+	setupPlatformSpecificMetrics(logger, processStats, systemMetrics, processMetrics)
 
 	return nil
 }
@@ -111,7 +137,7 @@ func infoReporter(serviceName, version string) func(_ monitoring.Mode, V monitor
 	}
 }
 
-func setupPlatformSpecificMetrics(logger *logp.Logger, processStats *process.Stats) {
+func setupPlatformSpecificMetrics(logger *logp.Logger, processStats *process.Stats, systemMetrics, processMetrics *monitoring.Registry) {
 	if isLinux() {
 		monitoring.NewFunc(processMetrics, "cgroup", InstanceCroupsReporter(logger, monitoringCgroupsHierarchyOverride), monitoring.Report)
 	}