refactor

Author: krishnamd-jkp

storage/grpc_client.go

@@ -127,7 +127,7 @@ type grpcStorageClient struct {
 
 	// configFeatureAttributes tracks client-level features that are enabled for this
 	// client instance, represented as an 8-bit bitmask.
-	configFeatureAttributes uint8
+	configFeatureAttributes uint32
 }
 
 func enableClientMetrics(ctx context.Context, s *settings, config storageConfig) (*metricsContext, error) {
@@ -251,13 +251,13 @@ func (c *grpcStorageClient) prepareDirectPathMetadata(ctx context.Context, targe
 	features |= c.configFeatureAttributes
 	// Merge all existing headers for this key from metadata.
 	for _, existing := range md[featureTrackerHeaderName] {
-		if decoded, err := decodeUint8(existing); err == nil {
+		if decoded, err := decodeUint32(existing); err == nil {
 			features |= decoded
 		}
 	}
 
 	if features > 0 {
-		md.Set(featureTrackerHeaderName, encodeUint8(features))
+		md.Set(featureTrackerHeaderName, encodeUint32(features))
 	}
 
 	return metadata.NewOutgoingContext(ctx, md), nil

storage/grpc_client_test.go

@@ -500,24 +500,24 @@ func TestPrepareDirectPathMetadata(t *testing.T) {
 func TestPrepareDirectPathMetadata_FeatureTracking(t *testing.T) {
 	tests := []struct {
 		desc            string
-		configFeatures  uint8
-		contextFeatures []featureCode
-		wantFeatures    uint8
+		configFeatures  uint32
+		contextFeatures []trackedFeature
+		wantFeatures    uint32
 	}{
 		{
 			desc:         "no features",
 			wantFeatures: 0,
 		},
 		{
 			desc:           "config features only",
-			configFeatures: uint8(featurePCU),
-			wantFeatures:   uint8(featurePCU),
+			configFeatures: uint32(1 << featurePCU),
+			wantFeatures:   uint32(1 << featurePCU),
 		},
 		{
 			desc:            "merged features",
-			configFeatures:  uint8(featurePCU),
-			contextFeatures: []featureCode{featureMultiStream},
-			wantFeatures:    uint8(featurePCU) | uint8(featureMultiStream),
+			configFeatures:  uint32(1 << featurePCU),
+			contextFeatures: []trackedFeature{featureMultistreamInMRD},
+			wantFeatures:    uint32(1<<featurePCU) | uint32(1<<featureMultistreamInMRD),
 		},
 	}
 
@@ -555,7 +555,7 @@ func TestPrepareDirectPathMetadata_FeatureTracking(t *testing.T) {
 				t.Fatalf("features header missing, want %d", tc.wantFeatures)
 			}
 
-			decoded, err := decodeUint8(got[0])
+			decoded, err := decodeUint32(got[0])
 			if err != nil {
 				t.Fatalf("failed to decode features: %v", err)
 			}

storage/reader.go

@@ -266,7 +266,7 @@ func (o *ObjectHandle) NewMultiRangeDownloader(ctx context.Context, opts ...MRDO
 		opt.apply(params)
 	}
 	if params.minConnections > 1 || params.maxConnections > 1 {
-		spanCtx = addFeatureAttributes(spanCtx, featureMultiStream)
+		spanCtx = addFeatureAttributes(spanCtx, featureMultistreamInMRD)
 	}
 	// This call will return the *MultiRangeDownloader with the .impl field set.
 	return o.c.tc.NewMultiRangeDownloader(spanCtx, params, storageOpts...)

storage/tracked_features.go

@@ -16,7 +16,7 @@ package storage
 
 // trackedFeature represents a specific client feature being tracked, represented
 // as a bit in a bitmask. Each feature corresponds to a specific bit position.
-type trackedFeature uint
+type trackedFeature uint32
 
 const (
 	featureMultistreamInMRD trackedFeature = 0

storage/tracker.go

@@ -16,71 +16,47 @@ package storage
 
 import (
 	"context"
-	"encoding/base64"
 
 	"github.com/googleapis/gax-go/v2/callctx"
 )
 
 const featureTrackerHeaderName = "x-goog-storage-go-features"
 
-// featureCode represents a specific client feature being tracked, represented
-// as a bit in a bitmask.
-type featureCode uint8
-
-const (
-	featureMultiStream featureCode = 1
-	featurePCU         featureCode = 2
-)
-
 // addFeatureAttributes adds the specified feature codes to the context.
 // Features are stored as a bitmask in the callctx headers and will be
 // injected into the outgoing request headers by the transport.
-func addFeatureAttributes(ctx context.Context, features ...featureCode) context.Context {
+func addFeatureAttributes(ctx context.Context, features ...trackedFeature) context.Context {
 	if len(features) == 0 {
 		return ctx
 	}
 
 	current := getFeatureAttributes(ctx)
 	updated := current
 	for _, f := range features {
-		updated |= uint8(f)
+		updated |= (1 << f)
 	}
 
 	if updated == current {
 		return ctx
 	}
 
-	return callctx.SetHeaders(ctx, featureTrackerHeaderName, encodeUint8(updated))
+	return callctx.SetHeaders(ctx, featureTrackerHeaderName, encodeUint32(uint32(updated)))
 }
 
 // getFeatureAttributes extracts and merges all feature attributes present in the context.
 // It returns a bitmask represented as a uint8.
-func getFeatureAttributes(ctx context.Context) uint8 {
+func getFeatureAttributes(ctx context.Context) uint32 {
 	ctxHeaders := callctx.HeadersFromContext(ctx)
 	if vals := ctxHeaders[featureTrackerHeaderName]; len(vals) > 0 {
 		// If multiple values are present in the context (e.g. from nested calls),
 		// merge them into a single bitmask.
-		var merged uint8
+		var merged uint32
 		for _, v := range vals {
-			if decoded, err := decodeUint8(v); err == nil {
+			if decoded, err := decodeUint32(v); err == nil {
 				merged |= decoded
 			}
 		}
 		return merged
 	}
 	return 0
 }
-
-// encodeUint8 encodes a uint8 bitmask into a base64 string for header injection.
-func encodeUint8(u uint8) string {
-	return base64.StdEncoding.EncodeToString([]byte{u})
-}
-
-// decodeUint8 decodes a base64 string back into a uint8 bitmask.
-func decodeUint8(s string) (uint8, error) {
-	b, err := base64.StdEncoding.DecodeString(s)
-	if err != nil || len(b) < 1 {
-		return 0, err
-	}
-	return b[0], nil
-}

storage/tracker_test.go

@@ -28,14 +28,14 @@ func TestAddFeatureAttributes(t *testing.T) {
 	}
 
 	// Add a single feature.
-	ctx = addFeatureAttributes(ctx, featureMultiStream)
-	if got := getFeatureAttributes(ctx); got != uint8(featureMultiStream) {
-		t.Errorf("getFeatureAttributes(MultiStream) = %d; want %d", got, featureMultiStream)
+	ctx = addFeatureAttributes(ctx, featureMultistreamInMRD)
+	if got := getFeatureAttributes(ctx); got != uint32(1<<featureMultistreamInMRD) {
+		t.Errorf("getFeatureAttributes(MultiStream) = %d; want %d", got, featureMultistreamInMRD)
 	}
 
 	// Add another feature (merge).
 	ctx = addFeatureAttributes(ctx, featurePCU)
-	want := uint8(featureMultiStream) | uint8(featurePCU)
+	want := uint32(1<<featureMultistreamInMRD) | uint32(1<<featurePCU)
 	if got := getFeatureAttributes(ctx); got != want {
 		t.Errorf("getFeatureAttributes(MultiStream | PCU) = %d; want %d", got, want)
 	}
@@ -46,19 +46,3 @@ func TestAddFeatureAttributes(t *testing.T) {
 		t.Errorf("getFeatureAttributes(MultiStream | PCU | PCU) = %d; want %d", got, want)
 	}
 }
-
-func TestEncodeDecodeUint8(t *testing.T) {
-	tests := []uint8{0, 1, 2, 127, 128, 255}
-
-	for _, want := range tests {
-		encoded := encodeUint8(want)
-		got, err := decodeUint8(encoded)
-		if err != nil {
-			t.Errorf("decodeUint8(%q) error: %v", encoded, err)
-			continue
-		}
-		if got != want {
-			t.Errorf("decodeUint8(encodeUint8(%d)) = %d; want %d", want, got, want)
-		}
-	}
-}