APIGOV-28984 Updates

pkg/agent/agent.go

@@ -53,7 +53,6 @@ type ConfigChangeHandler func()
 
 // ShutdownHandler - function that the agent may implement to be called when a shutdown request is received
 type ShutdownHandler func()
-
 type agentData struct {
 	agentResourceManager resource.Manager
 	teamJob              *centralTeamsCache
@@ -170,8 +169,6 @@ func InitializeWithAgentFeatures(centralCfg config.CentralConfig, agentFeaturesC
 	// call the metric services.
 	metricServicesConfigs := agentFeaturesCfg.GetMetricServicesConfigs()
 	agent.customUnitMetricServerManager = customunit.NewCustomUnitMetricServerManager(metricServicesConfigs, agent.cacheManager, centralCfg.GetAgentType())
-	ctx, ctxCancel := context.WithCancel(context.Background())
-	agent.customUnitMetricServerManager.HandleMetricReporting(ctx, ctxCancel)
 
 	if !agent.isInitialized {
 		err = handleInitialization()
@@ -436,6 +433,10 @@ func GetAuthProviderRegistry() oauth.ProviderRegistry {
 	return agent.authProviderRegistry
 }
 
+func GetCustomUnitMetricServerManager() *customunit.CustomUnitMetricServerManager {
+	return agent.customUnitMetricServerManager
+}
+
 // RegisterShutdownHandler - Registers shutdown handler
 func RegisterShutdownHandler(handler ShutdownHandler) {
 	agent.agentShutdownHandler = handler

pkg/customunit/client.go

@@ -90,23 +90,21 @@ type metricCollector interface {
 	AddCustomMetricDetail(models.CustomMetricDetail)
 }
 
-func (c *customUnitClient) MetricReporting() error {
+func (c *customUnitClient) MetricReporting() {
 	if err := c.createConnection(); err != nil {
 		//TODO:: Retry until the connection is stable
-		return err
+		return
 	}
 	c.metricReportingClient = cu.NewMetricReportingServiceClient(c.conn)
 	metricServiceInit := &cu.MetricServiceInit{}
 
 	client, err := c.metricReportingClient.MetricReporting(c.ctx, metricServiceInit, c.cOpts...)
 	if err != nil {
-		return err
+		return
 	}
 	c.metricReportingServiceClient = client
 	// process metrics
-	go c.processMetrics()
-	c.stopChan <- true
-	return nil
+	c.processMetrics()
 }
 
 // processMetrics will stream custom metrics
@@ -119,7 +117,7 @@ func (c *customUnitClient) processMetrics() {
 				c.cancelCtx()
 			}
 			c.MetricReporting()
-		case <-c.stopChan:
+		default:
 			metricReport, err := c.recv()
 			if err == io.EOF {
 				c.logger.Debug("stream finished")

pkg/customunit/client_test.go

@@ -73,8 +73,7 @@ func Test_MetricReporting(t *testing.T) {
 	ctx := context.Background()
 	fakeServer := &fakeCustomUnitMetricReportingServer{}
 	client, _ := createMRConnection(fakeServer, ctx)
-	err := client.MetricReporting()
-	fmt.Println(err)
+	client.MetricReporting()
 }
 
 func createMRConnection(fakeServer *fakeCustomUnitMetricReportingServer, ctx context.Context) (customUnitClient, error) {

pkg/customunit/manager.go

@@ -143,6 +143,6 @@ func (m *CustomUnitMetricServerManager) HandleMetricReporting(ctx context.Contex
 		factory := NewCustomUnitClientFactory(config.URL, m.cache, &customunits.QuotaInfo{})
 		client, _ := factory(ctx, cancelCtx)
 
-		client.MetricReporting()
+		go client.MetricReporting()
 	}
 }

pkg/transaction/metric/metricscollector.go

@@ -1,6 +1,7 @@
 package metric
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"os"
@@ -148,7 +149,8 @@ func GetMetricCollector() Collector {
 	if globalMetricCollector == nil && util.IsNotTest() {
 		globalMetricCollector = createMetricCollector()
 	}
-
+	ctx, ctxCancel := context.WithCancel(context.Background())
+	agent.GetCustomUnitMetricServerManager().HandleMetricReporting(ctx, ctxCancel)
 	return globalMetricCollector
 }