From 07361c70329fe79757221c312f062729e4b62446 Mon Sep 17 00:00:00 2001
From: Kaushik Surya <108111936+sky333999@users.noreply.github.com>
Date: Mon, 6 Jan 2025 23:49:31 -0600
Subject: [PATCH] Fix concurrent read write error for entity store (#1484)

---
 extension/entitystore/serviceprovider.go | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/extension/entitystore/serviceprovider.go b/extension/entitystore/serviceprovider.go
index b43fceeea5..5acee79a6a 100644
--- a/extension/entitystore/serviceprovider.go
+++ b/extension/entitystore/serviceprovider.go
@@ -64,6 +64,7 @@ type serviceprovider struct {
 	done             chan struct{}
 	logger           *zap.Logger
 	mutex            sync.RWMutex
+	logMutex         sync.RWMutex
 	// logFiles stores the service attributes that were configured for log files in CloudWatch Agent configuration.
 	// Example:
 	// "/opt/aws/amazon-cloudwatch-agent/logs/amazon-cloudwatch-agent.log": {ServiceName: "cloudwatch-agent"}
@@ -106,6 +107,8 @@ func (s *serviceprovider) getAutoScalingGroup() string {
 // addEntryForLogFile adds an association between a log file glob and a service attribute, as configured in the
 // CloudWatch Agent config.
 func (s *serviceprovider) addEntryForLogFile(logFileGlob LogFileGlob, serviceAttr ServiceAttribute) {
+	s.logMutex.Lock()
+	defer s.logMutex.Unlock()
 	if s.logFiles == nil {
 		s.logFiles = make(map[LogFileGlob]ServiceAttribute)
 	}
@@ -115,6 +118,8 @@ func (s *serviceprovider) addEntryForLogFile(logFileGlob LogFileGlob, serviceAtt
 // addEntryForLogGroup adds an association between a log group name and a service attribute, as observed from incoming
 // telemetry received by CloudWatch Agent.
 func (s *serviceprovider) addEntryForLogGroup(logGroupName LogGroupName, serviceAttr ServiceAttribute) {
+	s.logMutex.Lock()
+	defer s.logMutex.Unlock()
 	if s.logGroups == nil {
 		s.logGroups = make(map[LogGroupName]ServiceAttribute)
 	}
@@ -178,7 +183,8 @@ func (s *serviceprovider) serviceAttributeForLogGroup(logGroup LogGroupName) Ser
 	if logGroup == "" || s.logGroups == nil {
 		return ServiceAttribute{}
 	}
-
+	s.logMutex.RLock()
+	defer s.logMutex.RUnlock()
 	return s.logGroups[logGroup]
 }
 
@@ -186,7 +192,8 @@ func (s *serviceprovider) serviceAttributeForLogFile(logFile LogFileGlob) Servic
 	if logFile == "" || s.logFiles == nil {
 		return ServiceAttribute{}
 	}
-
+	s.logMutex.RLock()
+	defer s.logMutex.RUnlock()
 	return s.logFiles[logFile]
 }