Merge #158682

158682: cloud/azure: include env vars in client cache check r=dt a=dt

The tests of the azure client expect changes in environment variables to be reflected in subsequent uses of the client. However, the client was being cached without considering the env vars -- as they typically do not change during a process' lifetime outside of tests, violating the assumptions of these tests. Now the env var values are factored into the cache key.

Release note: none.
Epic: none.

Co-authored-by: David Taylor <[email protected]>

Author: craig[bot]

pkg/cloud/azure/azure_file_credential.go

@@ -62,6 +62,10 @@ func NewAzureFileCredential(
 	return c, nil
 }
 
+func credFileEnv() string {
+	return envutil.EnvOrDefaultString("COCKROACH_AZURE_APPLICATION_CREDENTIALS_FILE", "")
+}
+
 // NewDefaultAzureCredentialWithFile creates a credential chain that's simply
 // a FileCredential prepended to the chain in DefaultAzureCredential.
 func NewDefaultAzureCredentialWithFile(
@@ -74,7 +78,7 @@ func NewDefaultAzureCredentialWithFile(
 	var credentials []azcore.TokenCredential
 	var credErrors []string
 
-	credentialFileName := envutil.EnvOrDefaultString("COCKROACH_AZURE_APPLICATION_CREDENTIALS_FILE", "")
+	credentialFileName := credFileEnv()
 	fileCredentials, err := NewAzureFileCredential(credentialFileName, &FileCredentialOptions{ClientOptions: options.ClientOptions, testingKnobs: options.testingKnobs})
 	if err != nil {
 		credErrors = append(credErrors, errors.Wrap(err, "file credential").Error())

pkg/cloud/azure/azure_file_credentials_test.go

@@ -211,8 +211,7 @@ func TestAzureFileCredential(t *testing.T) {
 
 		t.Run("invalid-on-reload", func(t *testing.T) {
 			credFile := path.Join(tmpDir, "reload-on-error-invalid-on-reload")
-			cleanup := envutil.TestSetEnv(t, "COCKROACH_AZURE_APPLICATION_CREDENTIALS_FILE", credFile)
-			defer cleanup()
+			defer envutil.TestSetEnv(t, "COCKROACH_AZURE_APPLICATION_CREDENTIALS_FILE", credFile)()
 
 			require.NoError(t, writeAzureCredentialsFile(credFile, cfg.tenantID, cfg.clientID, cfg.clientSecret))
 

pkg/cloud/azure/azure_storage.go

@@ -28,6 +28,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/ioctx"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
@@ -222,10 +223,16 @@ type azureStorage struct {
 	settings  *cluster.Settings
 }
 
+type azCacheKey struct {
+	conf    cloudpb.ExternalStorage_Azure
+	envFile string
+	envID   string
+}
+
 var azClientCache struct {
 	syncutil.Mutex
 	// TODO(dt): make this an >1 item cache e.g. add a FIFO ring.
-	key    cloudpb.ExternalStorage_Azure
+	key    azCacheKey
 	set    time.Time
 	client *service.Client
 }
@@ -274,14 +281,19 @@ func makeAzureStorage(
 
 	var azClient *service.Client
 
-	clientConf := *conf
-	clientConf.Prefix = "" // Prefix is not part of the client identity.
+	clientID, _ := envutil.ExternalEnvString("AZURE_CLIENT_ID", 0)
+	cacheKey := azCacheKey{
+		conf:    *conf,
+		envFile: credFileEnv(),
+		envID:   clientID,
+	}
+	cacheKey.conf.Prefix = "" // Prefix is not part of the client identity.
 
 	if reuseSession.Get(&args.Settings.SV) {
 		func() {
 			azClientCache.Lock()
 			defer azClientCache.Unlock()
-			if cached := azClientCache.client; cached != nil && azClientCache.key == clientConf && timeutil.Since(azClientCache.set) < 2*time.Minute {
+			if cached := azClientCache.client; cached != nil && azClientCache.key == cacheKey && timeutil.Since(azClientCache.set) < 2*time.Minute {
 				azClient = cached
 			}
 		}()
@@ -342,7 +354,7 @@ func makeAzureStorage(
 	if reuseSession.Get(&args.Settings.SV) {
 		azClientCache.Lock()
 		defer azClientCache.Unlock()
-		azClientCache.key = clientConf
+		azClientCache.key = cacheKey
 		azClientCache.client = azClient
 		azClientCache.set = timeutil.Now()
 	}