feat(storage): add client feature tracking support (#14320)

Author: krishnamd-jkp

storage/grpc_client.go

@@ -124,6 +124,10 @@ type grpcStorageClient struct {
 	settings *settings
 	config   *storageConfig
 	dpDiag   string
+
+	// configFeatureAttributes tracks client-level features that are enabled for this
+	// client instance.
+	configFeatureAttributes uint32
 }
 
 func enableClientMetrics(ctx context.Context, s *settings, config storageConfig) (*metricsContext, error) {
@@ -240,6 +244,17 @@ func (c *grpcStorageClient) prepareDirectPathMetadata(ctx context.Context, targe
 			md.Set(requestParamsHeaderKey, reason)
 		}
 	}
+
+	// Client level feature tracking.
+	features := featureAttributes(ctx)
+	features |= c.configFeatureAttributes
+	// Merge all existing headers for this key from metadata.
+	features |= mergeFeatureAttributes(md[featureTrackerHeaderName])
+
+	if features > 0 {
+		md.Set(featureTrackerHeaderName, encodeUint32(features))
+	}
+
 	return metadata.NewOutgoingContext(ctx, md), nil
 }
 

storage/grpc_client_test.go

@@ -496,3 +496,73 @@ func TestPrepareDirectPathMetadata(t *testing.T) {
 		})
 	}
 }
+
+func TestPrepareDirectPathMetadata_FeatureTracking(t *testing.T) {
+	tests := []struct {
+		desc            string
+		configFeatures  uint32
+		contextFeatures []trackedFeature
+		wantFeatures    uint32
+	}{
+		{
+			desc:         "no features",
+			wantFeatures: 0,
+		},
+		{
+			desc:           "config features only",
+			configFeatures: uint32(1 << featurePCU),
+			wantFeatures:   uint32(1 << featurePCU),
+		},
+		{
+			desc:            "merged features",
+			configFeatures:  uint32(1 << featurePCU),
+			contextFeatures: []trackedFeature{featureMultistreamInMRD},
+			wantFeatures:    uint32(1<<featurePCU) | uint32(1<<featureMultistreamInMRD),
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.desc, func(t *testing.T) {
+			c := &grpcStorageClient{
+				config:                  &storageConfig{},
+				configFeatureAttributes: tc.configFeatures,
+			}
+
+			ctx := context.Background()
+			if len(tc.contextFeatures) > 0 {
+				ctx = addFeatureAttributes(ctx, tc.contextFeatures...)
+			}
+
+			newCtx, err := c.prepareDirectPathMetadata(ctx, directPathEndpointPrefix)
+			if err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+
+			md, ok := metadata.FromOutgoingContext(newCtx)
+			if !ok {
+				t.Fatal("metadata not found in context")
+			}
+
+			got := md.Get(featureTrackerHeaderName)
+			if tc.wantFeatures == 0 {
+				if len(got) > 0 {
+					t.Errorf("got features %q, want none", got[0])
+				}
+				return
+			}
+
+			if len(got) == 0 {
+				t.Fatalf("features header missing, want %d", tc.wantFeatures)
+			}
+
+			decoded, err := decodeUint32(got[0])
+			if err != nil {
+				t.Fatalf("failed to decode features: %v", err)
+			}
+
+			if decoded != tc.wantFeatures {
+				t.Errorf("got features %d, want %d", decoded, tc.wantFeatures)
+			}
+		})
+	}
+}

storage/http_client.go

@@ -45,6 +45,8 @@ import (
 
 // httpStorageClient is the HTTP-JSON API implementation of the transport-agnostic
 // storageClient interface.
+//
+// TODO(b/498422946): Add client feature tracker in HTTP client.
 type httpStorageClient struct {
 	creds                      *auth.Credentials
 	hc                         *http.Client

storage/pcu.go

@@ -179,6 +179,9 @@ func (w *Writer) initPCU(ctx context.Context) error {
 
 	s := newPCUSettings(cfg.MaxConcurrency)
 
+	// Track PCU operations using client feature tracking header.
+	ctx = addFeatureAttributes(ctx, featurePCU)
+
 	pCtx, cancel := context.WithCancel(ctx)
 
 	state := &pcuState{

storage/reader.go

@@ -265,7 +265,9 @@ func (o *ObjectHandle) NewMultiRangeDownloader(ctx context.Context, opts ...MRDO
 	for _, opt := range opts {
 		opt.apply(params)
 	}
-
+	if params.minConnections > 1 || params.maxConnections > 1 {
+		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

@@ -0,0 +1,63 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package storage
+
+import (
+	"context"
+
+	"github.com/googleapis/gax-go/v2/callctx"
+)
+
+const featureTrackerHeaderName = "x-goog-storage-go-features"
+
+// 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 ...trackedFeature) context.Context {
+	if len(features) == 0 {
+		return ctx
+	}
+
+	current := featureAttributes(ctx)
+	updated := current
+	for _, f := range features {
+		updated |= (1 << f)
+	}
+
+	if updated == current {
+		return ctx
+	}
+
+	return callctx.SetHeaders(ctx, featureTrackerHeaderName, encodeUint32(uint32(updated)))
+}
+
+// featureAttributes extracts and merges all feature attributes present in the context.
+// It returns a bitmask represented as a uint8.
+func featureAttributes(ctx context.Context) uint32 {
+	ctxHeaders := callctx.HeadersFromContext(ctx)
+	// If multiple values are present in the context (e.g. from nested calls),
+	// merge them into a single bitmask.
+	return mergeFeatureAttributes(ctxHeaders[featureTrackerHeaderName])
+}
+
+func mergeFeatureAttributes(vals []string) uint32 {
+	features := uint32(0)
+	for _, val := range vals {
+		if decoded, err := decodeUint32(val); err == nil {
+			features |= decoded
+		}
+	}
+	return features
+}

storage/tracker_test.go

@@ -0,0 +1,91 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package storage
+
+import (
+	"context"
+	"testing"
+)
+
+func TestAddFeatureAttributes(t *testing.T) {
+	ctx := context.Background()
+
+	// Initial features should be 0.
+	if got := featureAttributes(ctx); got != 0 {
+		t.Errorf("getFeatureAttributes(empty) = %d; want 0", got)
+	}
+
+	// Add a single feature.
+	ctx = addFeatureAttributes(ctx, featureMultistreamInMRD)
+	if got := featureAttributes(ctx); got != uint32(1<<featureMultistreamInMRD) {
+		t.Errorf("getFeatureAttributes(MultiStream) = %d; want %d", got, featureMultistreamInMRD)
+	}
+
+	// Add another feature (merge).
+	ctx = addFeatureAttributes(ctx, featurePCU)
+	want := uint32(1<<featureMultistreamInMRD) | uint32(1<<featurePCU)
+	if got := featureAttributes(ctx); got != want {
+		t.Errorf("getFeatureAttributes(MultiStream | PCU) = %d; want %d", got, want)
+	}
+
+	// Adding same feature should be idempotent.
+	ctx = addFeatureAttributes(ctx, featurePCU)
+	if got := featureAttributes(ctx); got != want {
+		t.Errorf("getFeatureAttributes(MultiStream | PCU | PCU) = %d; want %d", got, want)
+	}
+}
+
+func TestMergeFeatureAttributes(t *testing.T) {
+	tests := []struct {
+		name string
+		vals []string
+		want uint32
+	}{
+		{
+			name: "empty",
+			vals: []string{},
+			want: 0,
+		},
+		{
+			name: "single value",
+			vals: []string{encodeUint32(1)},
+			want: 1,
+		},
+		{
+			name: "multiple values",
+			vals: []string{encodeUint32(1), encodeUint32(2), encodeUint32(4)},
+			want: 7,
+		},
+		{
+			name: "overlapping values",
+			vals: []string{encodeUint32(3), encodeUint32(6)},
+			want: 7, // 011 | 110 = 111 (7)
+		},
+		{
+			name: "invalid values ignored",
+			vals: []string{encodeUint32(1), "invalid", encodeUint32(8)},
+			want: 9,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got := mergeFeatureAttributes(tc.vals)
+			if got != tc.want {
+				t.Errorf("mergeFeatureAttributes(%v) = %d; want %d", tc.vals, got, tc.want)
+			}
+		})
+	}
+}