refactor

pkg/storage/bucket/azure/bucket_client.go

@@ -22,16 +22,12 @@ func newBucketClient(cfg Config, name string, logger log.Logger, factory func(lo
 	bucketConfig.ContainerName = cfg.ContainerName
 	bucketConfig.MaxRetries = cfg.MaxRetries
 	bucketConfig.UserAssignedID = cfg.UserAssignedID
+	bucketConfig.HTTPConfig.Transport = cfg.Transport
 
 	if cfg.Endpoint != "" {
 		// azure.DefaultConfig has the default Endpoint, overwrite it only if a different one was explicitly provided.
 		bucketConfig.Endpoint = cfg.Endpoint
 	}
 
-	return factory(logger, bucketConfig, name, func(rt http.RoundTripper) http.RoundTripper {
-		if cfg.Transport != nil {
-			rt = cfg.Transport
-		}
-		return rt
-	})
+	return factory(logger, bucketConfig, name, nil)
 }

pkg/storage/bucket/object_client_adapter.go

@@ -2,16 +2,22 @@ package bucket
 
 import (
 	"context"
+	"fmt"
 	"io"
+	"net/http"
 	"slices"
 	"strings"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
 	"github.com/thanos-io/objstore"
 
 	"github.com/grafana/loki/v3/pkg/storage/chunk/client"
+	"github.com/grafana/loki/v3/pkg/storage/chunk/client/aws"
+	"github.com/grafana/loki/v3/pkg/storage/chunk/client/gcp"
+	"github.com/grafana/loki/v3/pkg/storage/chunk/client/hedging"
 )
 
 type ObjectClientAdapter struct {
@@ -21,9 +27,27 @@ type ObjectClientAdapter struct {
 	isRetryableErr       func(err error) bool
 }
 
-func NewObjectClientAdapter(bucket, hedgedBucket objstore.Bucket, logger log.Logger, opts ...ClientOptions) *ObjectClientAdapter {
-	if hedgedBucket == nil {
-		hedgedBucket = bucket
+func NewObjectClient(ctx context.Context, backend string, cfg Config, component string, hedgingCfg hedging.Config, disableRetries bool, logger log.Logger) (*ObjectClientAdapter, error) {
+	bucket, err := NewClient(ctx, backend, cfg, component, logger)
+	if err != nil {
+		return nil, fmt.Errorf("create bucket: %w", err)
+	}
+
+	hedgedBucket := bucket
+	if hedgingCfg.At != 0 {
+		hedgedTrasport, err := hedgingCfg.RoundTripperWithRegisterer(nil, prometheus.WrapRegistererWithPrefix("loki_", prometheus.DefaultRegisterer))
+		if err != nil {
+			return nil, fmt.Errorf("create hedged transport: %w", err)
+		}
+
+		if err := configureTransport(backend, hedgedTrasport, &cfg); err != nil {
+			return nil, err
+		}
+
+		bucket, err = NewClient(ctx, backend, cfg, component, logger)
+		if err != nil {
+			return nil, fmt.Errorf("create hedged bucket: %w", err)
+		}
 	}
 
 	o := &ObjectClientAdapter{
@@ -37,19 +61,33 @@ func NewObjectClientAdapter(bucket, hedgedBucket objstore.Bucket, logger log.Log
 		},
 	}
 
-	for _, opt := range opts {
-		opt(o)
+	switch backend {
+	case GCS:
+		o.isRetryableErr = gcp.IsRetryableErr
+	case S3:
+		o.isRetryableErr = aws.IsRetryableErr
 	}
 
-	return o
+	return o, nil
 }
 
-type ClientOptions func(*ObjectClientAdapter)
-
-func WithRetryableErrFunc(f func(err error) bool) ClientOptions {
-	return func(o *ObjectClientAdapter) {
-		o.isRetryableErr = f
+func configureTransport(backend string, rt http.RoundTripper, cfg *Config) error {
+	switch backend {
+	case S3:
+		cfg.S3.HTTP.Transport = rt
+	case GCS:
+		cfg.GCS.Transport = rt
+	case Azure:
+		cfg.Azure.Transport = rt
+	case Swift:
+		cfg.Swift.Transport = rt
+	case Filesystem:
+		// do nothing
+	default:
+		return fmt.Errorf("hedging not supported for backend: %s", backend)
 	}
+
+	return nil
 }
 
 func (o *ObjectClientAdapter) Stop() {

pkg/storage/bucket/object_client_adapter_test.go

@@ -6,10 +6,12 @@ import (
 	"sort"
 	"testing"
 
+	"github.com/go-kit/log"
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/loki/v3/pkg/storage/bucket/filesystem"
 	"github.com/grafana/loki/v3/pkg/storage/chunk/client"
+	"github.com/grafana/loki/v3/pkg/storage/chunk/client/hedging"
 )
 
 func TestObjectClientAdapter_List(t *testing.T) {
@@ -95,8 +97,12 @@ func TestObjectClientAdapter_List(t *testing.T) {
 		require.NoError(t, newBucket.Upload(context.Background(), "depply/nested/folder/b", buff))
 		require.NoError(t, newBucket.Upload(context.Background(), "depply/nested/folder/c", buff))
 
-		client := NewObjectClientAdapter(newBucket, nil, nil)
-		client.bucket = newBucket
+		client, err := NewObjectClient(context.Background(), "filesystem", Config{
+			StorageBackendConfig: StorageBackendConfig{
+				Filesystem: config,
+			},
+		}, "test", hedging.Config{}, false, log.NewNopLogger())
+		require.NoError(t, err)
 
 		storageObj, storageCommonPref, err := client.List(context.Background(), tt.prefix, tt.delimiter)
 		if tt.wantErr != nil {

pkg/storage/bucket/swift/bucket_client.go

@@ -4,8 +4,8 @@ import (
 	"github.com/go-kit/log"
 	"github.com/prometheus/common/model"
 	"github.com/thanos-io/objstore"
+	"github.com/thanos-io/objstore/exthttp"
 	"github.com/thanos-io/objstore/providers/swift"
-	yaml "gopkg.in/yaml.v2"
 )
 
 // NewBucketClient creates a new Swift bucket client
@@ -33,14 +33,9 @@ func NewBucketClient(cfg Config, _ string, logger log.Logger) (objstore.Bucket,
 		// Hard-coded defaults.
 		ChunkSize:              swift.DefaultConfig.ChunkSize,
 		UseDynamicLargeObjects: false,
+		HTTPConfig:             exthttp.DefaultHTTPConfig,
 	}
+	bucketConfig.HTTPConfig.Transport = cfg.Transport
 
-	// Thanos currently doesn't support passing the config as is, but expects a YAML,
-	// so we're going to serialize it.
-	serialized, err := yaml.Marshal(bucketConfig)
-	if err != nil {
-		return nil, err
-	}
-
-	return swift.NewContainer(logger, serialized, nil)
+	return swift.NewContainerFromConfig(logger, &bucketConfig, false, nil)
 }

pkg/storage/bucket/swift/config.go

@@ -2,6 +2,7 @@ package swift
 
 import (
 	"flag"
+	"net/http"
 	"time"
 )
 
@@ -26,6 +27,9 @@ type Config struct {
 	MaxRetries        int           `yaml:"max_retries"`
 	ConnectTimeout    time.Duration `yaml:"connect_timeout"`
 	RequestTimeout    time.Duration `yaml:"request_timeout"`
+
+	// Allow upstream callers to inject a round tripper
+	Transport http.RoundTripper `yaml:"-"`
 }
 
 // RegisterFlags registers the flags for Swift storage

pkg/storage/factory.go

@@ -615,12 +615,16 @@ func (c *ClientMetrics) Unregister() {
 
 // NewObjectClient makes a new StorageClient with the prefix in the front.
 func NewObjectClient(name, component string, cfg Config, clientMetrics ClientMetrics) (client.ObjectClient, error) {
+	if cfg.UseThanosObjstore {
+		return bucket.NewObjectClient(context.Background(), name, cfg.ObjectStore, component, cfg.Hedging, cfg.CongestionControl.Enabled, util_log.Logger)
+	}
+
 	actual, err := internalNewObjectClient(name, component, cfg, clientMetrics)
 	if err != nil {
 		return nil, err
 	}
 
-	if cfg.UseThanosObjstore || cfg.ObjectPrefix == "" {
+	if cfg.ObjectPrefix == "" {
 		return actual, nil
 	} else {
 		prefix := strings.Trim(cfg.ObjectPrefix, "/") + "/"
@@ -659,9 +663,6 @@ func internalNewObjectClient(storeName, component string, cfg Config, clientMetr
 			s3Cfg.BackoffConfig.MaxRetries = 1
 		}
 
-		if cfg.UseThanosObjstore {
-			return aws.NewS3ThanosObjectClient(context.Background(), cfg.ObjectStore, component, util_log.Logger, cfg.Hedging)
-		}
 		return aws.NewS3ObjectClient(s3Cfg, cfg.Hedging)
 
 	case types.StorageTypeAlibabaCloud:
@@ -691,9 +692,6 @@ func internalNewObjectClient(storeName, component string, cfg Config, clientMetr
 		if cfg.CongestionControl.Enabled {
 			gcsCfg.EnableRetries = false
 		}
-		if cfg.UseThanosObjstore {
-			return gcp.NewGCSThanosObjectClient(context.Background(), cfg.ObjectStore, component, util_log.Logger, cfg.Hedging)
-		}
 		return gcp.NewGCSObjectClient(context.Background(), gcsCfg, cfg.Hedging)
 
 	case types.StorageTypeAzure:
@@ -705,9 +703,6 @@ func internalNewObjectClient(storeName, component string, cfg Config, clientMetr
 			}
 			azureCfg = (azure.BlobStorageConfig)(nsCfg)
 		}
-		if cfg.UseThanosObjstore {
-			return azure.NewBlobStorageThanosObjectClient(context.Background(), cfg.ObjectStore, component, util_log.Logger, cfg.Hedging)
-		}
 		return azure.NewBlobStorage(&azureCfg, clientMetrics.AzureMetrics, cfg.Hedging)
 
 	case types.StorageTypeSwift:
@@ -757,10 +752,6 @@ func internalNewObjectClient(storeName, component string, cfg Config, clientMetr
 		return ibmcloud.NewCOSObjectClient(cosCfg, cfg.Hedging)
 
 	default:
-		if cfg.UseThanosObjstore {
-			return nil, fmt.Errorf("Unrecognized storage client %v, choose one of: %s", storeName, strings.Join(cfg.ObjectStore.SupportedBackends(), ", "))
-		}
-
 		return nil, fmt.Errorf("Unrecognized storage client %v, choose one of: %v, %v, %v, %v, %v, %v, %v, %v, %v", storeName, types.StorageTypeAWS, types.StorageTypeS3, types.StorageTypeGCS, types.StorageTypeAzure, types.StorageTypeAlibabaCloud, types.StorageTypeSwift, types.StorageTypeBOS, types.StorageTypeCOS, types.StorageTypeFileSystem)
 	}
 }