Fix CachingFactory concurrent usage issues (#1707)

Signed-off-by: andriikushch <[email protected]>

Author: andriikushch

go.mod

@@ -27,6 +27,7 @@ require (
 	github.com/r3labs/diff/v3 v3.0.1
 	github.com/stretchr/testify v1.10.0
 	github.com/urfave/cli/v2 v2.27.7
+	go.uber.org/atomic v1.11.0
 	golang.org/x/exp v0.0.0-20240823005443-9b4947da3948
 	golang.org/x/sync v0.15.0
 	gopkg.in/yaml.v2 v2.4.0

go.sum

@@ -111,6 +111,8 @@ github.com/xhit/go-str2duration/v2 v2.1.0 h1:lxklc02Drh6ynqX+DdPyp5pCKLUQpRT8bp8
 github.com/xhit/go-str2duration/v2 v2.1.0/go.mod h1:ohY8p+0f07DiV6Em5LKB0s2YpLtXVyJfNt1+BlmyAsU=
 github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1 h1:gEOO8jv9F4OT7lGCjxCBTO/36wtF6j2nSip77qHd4x4=
 github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1/go.mod h1:Ohn+xnUBiLI6FVj/9LpzZWtj1/D6lUovWYBkxHVV3aM=
+go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE=
+go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
 golang.org/x/exp v0.0.0-20240823005443-9b4947da3948 h1:kx6Ds3MlpiUHKj7syVnbp57++8WpuKPcR5yjLBjvLEA=
 golang.org/x/exp v0.0.0-20240823005443-9b4947da3948/go.mod h1:akd2r19cwCdwSwWeIdzYQGa/EZZyqcOdwWiwj5L5eKQ=
 golang.org/x/sync v0.15.0 h1:KWH3jNZsfyT6xfAfKiz6MRNmd46ByHDYaZ7KSkCtdW8=

pkg/clients/v1/factory.go

@@ -46,6 +46,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/storagegateway/storagegatewayiface"
 	"github.com/aws/aws-sdk-go/service/sts"
 	"github.com/aws/aws-sdk-go/service/sts/stsiface"
+	"go.uber.org/atomic"
 
 	"github.com/prometheus-community/yet-another-cloudwatch-exporter/pkg/clients"
 	"github.com/prometheus-community/yet-another-cloudwatch-exporter/pkg/clients/account"
@@ -64,8 +65,8 @@ type CachingFactory struct {
 	stscache         map[model.Role]stsiface.STSAPI
 	iamcache         map[model.Role]iamiface.IAMAPI
 	clients          map[model.Role]map[string]*cachedClients
-	cleared          bool
-	refreshed        bool
+	cleared          *atomic.Bool
+	refreshed        *atomic.Bool
 	mu               sync.Mutex
 	fips             bool
 	logger           *slog.Logger
@@ -175,16 +176,21 @@ func NewFactory(logger *slog.Logger, jobsCfg model.JobsConfig, fips bool) *Cachi
 		iamcache:         iamcache,
 		clients:          cache,
 		fips:             fips,
-		cleared:          false,
-		refreshed:        false,
+		cleared:          atomic.NewBool(false),
+		refreshed:        atomic.NewBool(false),
 		logger:           logger,
 	}
 }
 
-// Refresh and Clear help to avoid using lock primitives by asserting that
-// there are no ongoing writes to the map.
 func (c *CachingFactory) Clear() {
-	if c.cleared {
+	if c.cleared.Load() {
+		return
+	}
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.cleared.Load() {
 		return
 	}
 
@@ -204,19 +210,19 @@ func (c *CachingFactory) Clear() {
 			cachedClient.tagging = nil
 		}
 	}
-	c.cleared = true
-	c.refreshed = false
+	c.cleared.Store(true)
+	c.refreshed.Store(false)
 }
 
 func (c *CachingFactory) Refresh() {
-	if c.refreshed {
+	if c.refreshed.Load() {
 		return
 	}
 
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	// Double check Refresh wasn't called concurrently
-	if c.refreshed {
+	if c.refreshed.Load() {
 		return
 	}
 
@@ -248,8 +254,8 @@ func (c *CachingFactory) Refresh() {
 		}
 	}
 
-	c.cleared = false
-	c.refreshed = true
+	c.cleared.Store(false)
+	c.refreshed.Store(true)
 }
 
 func createCloudWatchClient(logger *slog.Logger, s *session.Session, region *string, role model.Role, fips bool) cloudwatch_client.Client {
@@ -282,7 +288,7 @@ func createAccountClient(logger *slog.Logger, sts stsiface.STSAPI, iam iamiface.
 }
 
 func (c *CachingFactory) GetCloudwatchClient(region string, role model.Role, concurrency cloudwatch_client.ConcurrencyConfig) cloudwatch_client.Client {
-	if !c.refreshed {
+	if !c.refreshed.Load() {
 		// if we have not refreshed then we need to lock in case we are accessing concurrently
 		c.mu.Lock()
 		defer c.mu.Unlock()
@@ -295,7 +301,7 @@ func (c *CachingFactory) GetCloudwatchClient(region string, role model.Role, con
 }
 
 func (c *CachingFactory) GetTaggingClient(region string, role model.Role, concurrencyLimit int) tagging.Client {
-	if !c.refreshed {
+	if !c.refreshed.Load() {
 		// if we have not refreshed then we need to lock in case we are accessing concurrently
 		c.mu.Lock()
 		defer c.mu.Unlock()
@@ -308,7 +314,7 @@ func (c *CachingFactory) GetTaggingClient(region string, role model.Role, concur
 }
 
 func (c *CachingFactory) GetAccountClient(region string, role model.Role) account.Client {
-	if !c.refreshed {
+	if !c.refreshed.Load() {
 		// if we have not refreshed then we need to lock in case we are accessing concurrently
 		c.mu.Lock()
 		defer c.mu.Unlock()

pkg/clients/v1/factory_test.go

@@ -16,6 +16,7 @@ import (
 	"fmt"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
@@ -25,6 +26,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/sts/stsiface"
 	"github.com/prometheus/common/promslog"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
 
 	"github.com/prometheus-community/yet-another-cloudwatch-exporter/pkg/clients/cloudwatch"
 	"github.com/prometheus-community/yet-another-cloudwatch-exporter/pkg/model"
@@ -82,15 +84,21 @@ func TestNewClientCache(t *testing.T) {
 			"an empty config gives an empty cache",
 			model.JobsConfig{},
 			false,
-			&CachingFactory{logger: promslog.NewNopLogger()},
+			&CachingFactory{
+				logger:    promslog.NewNopLogger(),
+				refreshed: atomic.NewBool(false),
+				cleared:   atomic.NewBool(false),
+			},
 		},
 		{
 			"if fips is set then the clients has fips",
 			model.JobsConfig{},
 			true,
 			&CachingFactory{
-				fips:   true,
-				logger: promslog.NewNopLogger(),
+				fips:      true,
+				logger:    promslog.NewNopLogger(),
+				refreshed: atomic.NewBool(false),
+				cleared:   atomic.NewBool(false),
 			},
 		},
 		{
@@ -153,7 +161,9 @@ func TestNewClientCache(t *testing.T) {
 						"ap-northeast-3": &cachedClients{},
 					},
 				},
-				logger: promslog.NewNopLogger(),
+				logger:    promslog.NewNopLogger(),
+				refreshed: atomic.NewBool(false),
+				cleared:   atomic.NewBool(false),
 			},
 		},
 		{
@@ -239,7 +249,9 @@ func TestNewClientCache(t *testing.T) {
 						"ap-northeast-1": &cachedClients{onlyStatic: true},
 					},
 				},
-				logger: promslog.NewNopLogger(),
+				logger:    promslog.NewNopLogger(),
+				refreshed: atomic.NewBool(false),
+				cleared:   atomic.NewBool(false),
 			},
 		},
 		{
@@ -362,7 +374,9 @@ func TestNewClientCache(t *testing.T) {
 						"ap-northeast-3": &cachedClients{},
 					},
 				},
-				logger: promslog.NewNopLogger(),
+				logger:    promslog.NewNopLogger(),
+				refreshed: atomic.NewBool(false),
+				cleared:   atomic.NewBool(false),
 			},
 		},
 		{
@@ -451,7 +465,9 @@ func TestNewClientCache(t *testing.T) {
 						"ap-northeast-1": &cachedClients{onlyStatic: true},
 					},
 				},
-				logger: promslog.NewNopLogger(),
+				logger:    promslog.NewNopLogger(),
+				refreshed: atomic.NewBool(false),
+				cleared:   atomic.NewBool(false),
 			},
 		},
 	}
@@ -463,12 +479,12 @@ func TestNewClientCache(t *testing.T) {
 			cache := NewFactory(promslog.NewNopLogger(), test.jobsCfg, test.fips)
 			t.Logf("the cache is: %v", cache)
 
-			if test.cache.cleared != cache.cleared {
+			if test.cache.cleared.Load() != cache.cleared.Load() {
 				t.Logf("`cleared` not equal got %v, expected %v", cache.cleared, test.cache.cleared)
 				t.Fail()
 			}
 
-			if test.cache.refreshed != cache.refreshed {
+			if test.cache.refreshed.Load() != cache.refreshed.Load() {
 				t.Logf("`refreshed` not equal got %v, expected %v", cache.refreshed, test.cache.refreshed)
 				t.Fail()
 			}
@@ -497,7 +513,7 @@ func TestClear(t *testing.T) {
 			"a new clear clears all clients",
 			&CachingFactory{
 				session: mock.Session,
-				cleared: false,
+				cleared: atomic.NewBool(false),
 				mu:      sync.Mutex{},
 				stscache: map[model.Role]stsiface.STSAPI{
 					{}: nil,
@@ -512,13 +528,14 @@ func TestClear(t *testing.T) {
 						},
 					},
 				},
-				logger: promslog.NewNopLogger(),
+				logger:    promslog.NewNopLogger(),
+				refreshed: atomic.NewBool(false),
 			},
 		},
 		{
 			"A second call to clear does nothing",
 			&CachingFactory{
-				cleared: true,
+				cleared: atomic.NewBool(true),
 				mu:      sync.Mutex{},
 				session: mock.Session,
 				stscache: map[model.Role]stsiface.STSAPI{
@@ -533,7 +550,8 @@ func TestClear(t *testing.T) {
 						},
 					},
 				},
-				logger: promslog.NewNopLogger(),
+				logger:    promslog.NewNopLogger(),
+				refreshed: atomic.NewBool(false),
 			},
 		},
 	}
@@ -542,11 +560,11 @@ func TestClear(t *testing.T) {
 		test := l
 		t.Run(test.description, func(t *testing.T) {
 			test.cache.Clear()
-			if !test.cache.cleared {
+			if !test.cache.cleared.Load() {
 				t.Log("Cache cleared flag not set")
 				t.Fail()
 			}
-			if test.cache.refreshed {
+			if test.cache.refreshed.Load() {
 				t.Log("Cache cleared flag set")
 				t.Fail()
 			}
@@ -591,7 +609,8 @@ func TestRefresh(t *testing.T) {
 			"a new refresh creates clients",
 			&CachingFactory{
 				session:   mock.Session,
-				refreshed: false,
+				refreshed: atomic.NewBool(false),
+				cleared:   atomic.NewBool(false),
 				mu:        sync.Mutex{},
 				stscache: map[model.Role]stsiface.STSAPI{
 					{}: nil,
@@ -613,7 +632,8 @@ func TestRefresh(t *testing.T) {
 			"a new refresh with static only creates only cloudwatch",
 			&CachingFactory{
 				session:   mock.Session,
-				refreshed: false,
+				refreshed: atomic.NewBool(false),
+				cleared:   atomic.NewBool(false),
 				mu:        sync.Mutex{},
 				stscache: map[model.Role]stsiface.STSAPI{
 					{}: nil,
@@ -635,7 +655,8 @@ func TestRefresh(t *testing.T) {
 		{
 			"A second call to refreshed does nothing",
 			&CachingFactory{
-				refreshed: true,
+				refreshed: atomic.NewBool(true),
+				cleared:   atomic.NewBool(false),
 				mu:        sync.Mutex{},
 				session:   mock.Session,
 				stscache: map[model.Role]stsiface.STSAPI{
@@ -662,12 +683,12 @@ func TestRefresh(t *testing.T) {
 			t.Parallel()
 			test.cache.Refresh()
 
-			if !test.cache.refreshed {
+			if !test.cache.refreshed.Load() {
 				t.Log("Cache refreshed flag not set")
 				t.Fail()
 			}
 
-			if test.cache.cleared {
+			if test.cache.cleared.Load() {
 				t.Log("Cache cleared flag set")
 				t.Fail()
 			}
@@ -753,7 +774,7 @@ func testGetAWSClient(
 		{
 			"locks during unrefreshed parallel call",
 			&CachingFactory{
-				refreshed: false,
+				refreshed: atomic.NewBool(false),
 				mu:        sync.Mutex{},
 				session:   mock.Session,
 				stscache: map[model.Role]stsiface.STSAPI{
@@ -775,7 +796,7 @@ func testGetAWSClient(
 		{
 			"returns clients if available",
 			&CachingFactory{
-				refreshed: true,
+				refreshed: atomic.NewBool(true),
 				session:   mock.Session,
 				mu:        sync.Mutex{},
 				stscache: map[model.Role]stsiface.STSAPI{
@@ -797,7 +818,7 @@ func testGetAWSClient(
 		{
 			"creates a new clients if not available",
 			&CachingFactory{
-				refreshed: true,
+				refreshed: atomic.NewBool(true),
 				session:   mock.Session,
 				mu:        sync.Mutex{},
 				stscache: map[model.Role]stsiface.STSAPI{
@@ -1150,6 +1171,45 @@ func TestSTSResolvesFIPSEnabledEndpoints(t *testing.T) {
 	}
 }
 
+func TestRaceConditionRefreshClear(t *testing.T) {
+	t.Parallel()
+
+	// Create a factory with the test config
+	factory := NewFactory(promslog.NewNopLogger(), model.JobsConfig{}, false)
+
+	// Number of concurrent operations to perform
+	iterations := 100
+
+	// Use WaitGroup to synchronize goroutines
+	var wg sync.WaitGroup
+	wg.Add(iterations) // For both Refresh and Clear calls
+
+	// Start function to run concurrent operations
+	for i := 0; i < iterations; i++ {
+		// Launch goroutine to call Refresh
+		go func() {
+			defer wg.Done()
+			factory.Refresh()
+			factory.Clear()
+		}()
+	}
+
+	// Create a channel to signal completion
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	// Wait for either completion or timeout
+	select {
+	case <-done:
+		// Test completed successfully
+	case <-time.After(60 * time.Second):
+		require.Fail(t, "Test timed out after 60 seconds")
+	}
+}
+
 func testAWSClient(
 	t *testing.T,
 	name string,