cloud/azure: add optional azure client caching

This adds a signle-slot, process-wide in-memory cache of the most recently used
azure SDK storage client allowing subsequent attempts to open clients with the
same configuration to reuse the already open client. This can allow the SDK's
internal caching of things like IMDS tokens to be used across separate CRDB
operations on separate files which currently each open their own client.

This can be particularly impactful for avoiding hitting the IMDS rate limit.

Release note: none.
Epic: none.

Author: dt

pkg/cloud/azure/BUILD.bazel

@@ -24,6 +24,7 @@ go_library(
         "//pkg/util/envutil",
         "//pkg/util/ioctx",
         "//pkg/util/log",
+        "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "//pkg/util/tracing",
         "@com_github_azure_azure_sdk_for_go_sdk_azcore//:azcore",

pkg/cloud/azure/azure_storage.go

@@ -29,6 +29,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/util/ioctx"
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/errors"
@@ -55,6 +56,13 @@ var tryTimeout = settings.RegisterDurationSetting(
 	"the timeout for individual retry attempts in Azure operations",
 	60*time.Second)
 
+var reuseSession = settings.RegisterBoolSetting(
+	settings.ApplicationLevel,
+	"cloudstorage.azure.session_reuse.enabled",
+	"persist the last opened azure client and re-use it when opening a new client with the same argument (some settings may take 2mins to take effect)",
+	false,
+)
+
 // A note on Azure authentication:
 //
 // The standardized way to authenticate a third-party identity to the Azure
@@ -214,6 +222,14 @@ type azureStorage struct {
 	settings  *cluster.Settings
 }
 
+var azClientCache struct {
+	syncutil.Mutex
+	// TODO(dt): make this an >1 item cache e.g. add a FIFO ring.
+	key    cloudpb.ExternalStorage_Azure
+	set    time.Time
+	client *service.Client
+}
+
 var _ cloud.ExternalStorage = &azureStorage{}
 
 func makeAzureStorage(
@@ -224,6 +240,7 @@ func makeAzureStorage(
 	if conf == nil {
 		return nil, errors.Errorf("azure upload requested but info missing")
 	}
+
 	env, err := azure.EnvironmentFromName(conf.Environment)
 	if err != nil {
 		return nil, errors.Wrap(err, "azure environment")
@@ -256,6 +273,29 @@ func makeAzureStorage(
 	opts.Retry.TryTimeout = tryTimeout.Get(&args.Settings.SV)
 
 	var azClient *service.Client
+
+	clientConf := *conf
+	clientConf.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 {
+				azClient = cached
+			}
+		}()
+		if azClient != nil {
+			return &azureStorage{
+				conf:      conf,
+				ioConf:    args.IOConf,
+				container: azClient.NewContainerClient(conf.Container),
+				prefix:    conf.Prefix,
+				settings:  args.Settings,
+			}, nil
+		}
+	}
+
 	switch conf.Auth {
 	case cloudpb.AzureAuth_LEGACY:
 		credential, err := azblob.NewSharedKeyCredential(conf.AccountName, conf.AccountKey)
@@ -299,6 +339,14 @@ func makeAzureStorage(
 		return nil, errors.Errorf("unsupported value %s for %s", conf.Auth, cloud.AuthParam)
 	}
 
+	if reuseSession.Get(&args.Settings.SV) {
+		azClientCache.Lock()
+		defer azClientCache.Unlock()
+		azClientCache.key = clientConf
+		azClientCache.client = azClient
+		azClientCache.set = timeutil.Now()
+	}
+
 	return &azureStorage{
 		conf:      conf,
 		ioConf:    args.IOConf,