replace mutex+map with sync.Map

Signed-off-by: Etai Lev Ran <[email protected]>

Author: elevran

pkg/epp/datalayer/attributemap.go

@@ -36,59 +36,54 @@ type AttributeMap interface {
 
 // Attributes provides a goroutine safe implementation of AttributeMap.
 type Attributes struct {
-	mu   sync.RWMutex
-	data map[string]Cloneable
+	data sync.Map
 }
 
 // NewAttributes return a new attribute map instance.
 func NewAttributes() *Attributes {
 	return &Attributes{
-		data: make(map[string]Cloneable),
+		data: sync.Map{},
 	}
 }
 
 // Put adds (or updates) an attribute in the map.
 func (a *Attributes) Put(key string, value Cloneable) {
-	a.mu.Lock()
-	defer a.mu.Unlock()
-
-	a.data[key] = value // TODO: Clone into map?
+	a.data.Store(key, value) // TODO: Clone into map?
 }
 
 // Get returns an attribute from the map.
 func (a *Attributes) Get(key string) (Cloneable, bool) {
-	a.mu.RLock()
-	defer a.mu.RUnlock()
-
-	val, ok := a.data[key]
+	val, ok := a.data.Load(key)
 	if !ok {
 		return nil, false
 	}
-	return val.Clone(), true
+	if cloneable, ok := val.(Cloneable); ok {
+		return cloneable.Clone(), true
+	}
+	return nil, false // shouldn't happen since Put accepts Cloneables only
 }
 
 // Keys returns an array of all the names of attributes stored in the map.
 func (a *Attributes) Keys() []string {
-	a.mu.RLock()
-	defer a.mu.RUnlock()
-
-	keys := make([]string, 0, len(a.data))
-	for k := range a.data {
-		keys = append(keys, k)
-	}
+	keys := []string{}
+	a.data.Range(func(key, _ any) bool {
+		if k, ok := key.(string); ok {
+			keys = append(keys, k)
+		}
+		return true // continue iteration
+	})
 	return keys
 }
 
 // Clone the attributes object itself.
 func (a *Attributes) Clone() *Attributes {
-	a.mu.RLock()
-	defer a.mu.RUnlock()
-
-	m := make(map[string]Cloneable, len(a.data))
-	for k, v := range a.data {
-		m[k] = v.Clone()
-	}
-	return &Attributes{
-		data: m,
+	cloned := &Attributes{
+		data: sync.Map{},
 	}
+
+	a.data.Range(func(k, v interface{}) bool {
+		cloned.data.Store(k, v)
+		return true
+	})
+	return cloned
 }

pkg/epp/datalayer/datasource.go

@@ -43,11 +43,6 @@ type DataSource interface {
 	// the data source output type upon registration.
 	AddExtractor(extractor Extractor) error
 
-	// TODO: the following is useful for a data source that operates on
-	// endpoints and might not be relevant for "global/system" collectors which
-	// might not need the concept of an endpoint. This can be split, if needed,
-	// to a separate interface in the future.
-
 	// AddEndpoint adds an endpoint to collect from.
 	AddEndpoint(ep Endpoint) error
 
@@ -76,8 +71,7 @@ var (
 // DataSourceRegistry stores named data sources and makes them
 // accessible to GIE subsystems.
 type DataSourceRegistry struct {
-	mu      sync.RWMutex
-	sources map[string]DataSource
+	sources sync.Map
 }
 
 // Register adds a source to the registry.
@@ -86,13 +80,10 @@ func (dsr *DataSourceRegistry) Register(src DataSource) error {
 		return errors.New("unable to register a nil data source")
 	}
 
-	dsr.mu.Lock()
-	defer dsr.mu.Unlock()
-
-	if _, found := dsr.sources[src.Name()]; found {
+	if _, found := dsr.sources.Load(src.Name()); found {
 		return fmt.Errorf("unable to register duplicate data source: %s", src.Name())
 	}
-	dsr.sources[src.Name()] = src
+	dsr.sources.Store(src.Name(), src)
 	return nil
 }
 
@@ -102,10 +93,10 @@ func (dsr *DataSourceRegistry) GetNamedSource(name string) (DataSource, bool) {
 		return nil, false
 	}
 
-	dsr.mu.RLock()
-	defer dsr.mu.RUnlock()
-	if ds, found := dsr.sources[name]; found {
-		return ds, true
+	if val, found := dsr.sources.Load(name); found {
+		if ds, ok := val.(DataSource); ok {
+			return ds, true
+		} // ignore type assertion failures and fall through
 	}
 	return nil, false
 }
@@ -122,15 +113,15 @@ func (dsr *DataSourceRegistry) AddEndpoint(ep Endpoint) error {
 		return nil
 	}
 
-	dsr.mu.RLock()
-	defer dsr.mu.RUnlock()
 	errs := []error{}
-
-	for _, ds := range dsr.sources {
-		if err := ds.AddEndpoint(ep); err != nil {
-			errs = append(errs, err)
+	dsr.sources.Range(func(_, val interface{}) bool {
+		if ds, ok := val.(DataSource); ok {
+			if err := ds.AddEndpoint(ep); err != nil {
+				errs = append(errs, err)
+			}
 		}
-	}
+		return true
+	})
 	return errors.Join(errs...)
 }
 
@@ -142,15 +133,15 @@ func (dsr *DataSourceRegistry) RemoveEndpoint(ep Endpoint) error {
 		return nil
 	}
 
-	dsr.mu.RLock()
-	defer dsr.mu.RUnlock()
 	errs := []error{}
-
-	for _, ds := range dsr.sources {
-		if err := ds.RemoveEndpoint(ep); err != nil {
-			errs = append(errs, err)
+	dsr.sources.Range(func(_, val interface{}) bool {
+		if ds, ok := val.(DataSource); ok {
+			if err := ds.RemoveEndpoint(ep); err != nil {
+				errs = append(errs, err)
+			}
 		}
-	}
+		return true
+	})
 	return errors.Join(errs...)
 }