feat(v2): add TransportTelemetryData for dynamic transport attributes (#481)

westarle · web-flow · commit 8a7caf0014c9 · 2026-03-12T13:19:03.000-07:00
diff --git a/v2/telemetry.go b/v2/telemetry.go
@@ -30,11 +30,56 @@
 package gax
 
 import (
+	"context"
+
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
 )
 
+// TransportTelemetryData contains mutable telemetry information that the transport
+// layer (e.g. gRPC or HTTP) populates during an RPC. This allows gax.Invoke to
+// correctly emit metric data without directly importing those transport layers.
+// This is an EXPERIMENTAL struct and should not be used by external consumers.
+type TransportTelemetryData struct {
+	serverAddress      string
+	serverPort         int
+	responseStatusCode int
+}
+
+// SetServerAddress sets the server address.
+func (d *TransportTelemetryData) SetServerAddress(addr string) { d.serverAddress = addr }
+
+// ServerAddress returns the server address.
+func (d *TransportTelemetryData) ServerAddress() string { return d.serverAddress }
+
+// SetServerPort sets the server port.
+func (d *TransportTelemetryData) SetServerPort(port int) { d.serverPort = port }
+
+// ServerPort returns the server port.
+func (d *TransportTelemetryData) ServerPort() int { return d.serverPort }
+
+// SetResponseStatusCode sets the response status code.
+func (d *TransportTelemetryData) SetResponseStatusCode(code int) { d.responseStatusCode = code }
+
+// ResponseStatusCode returns the response status code.
+func (d *TransportTelemetryData) ResponseStatusCode() int { return d.responseStatusCode }
+
+// transportTelemetryKey is the private context key used to inject TransportTelemetryData
+type transportTelemetryKey struct{}
+
+// InjectTransportTelemetry injects a mutable TransportTelemetryData pointer into the context.
+// Experimental: This function is subject to breaking changes.
+func InjectTransportTelemetry(ctx context.Context, data *TransportTelemetryData) context.Context {
+	return context.WithValue(ctx, transportTelemetryKey{}, data)
+}
+
+// ExtractTransportTelemetry retrieves a mutable TransportTelemetryData pointer from the context.
+func ExtractTransportTelemetry(ctx context.Context) *TransportTelemetryData {
+	data, _ := ctx.Value(transportTelemetryKey{}).(*TransportTelemetryData)
+	return data
+}
+
 const (
 	metricName        = "gcp.client.request.duration"
 	metricDescription = "Duration of the request to the Google Cloud API"
diff --git a/v2/telemetry_test.go b/v2/telemetry_test.go
@@ -234,3 +234,29 @@ func TestNoSDKImport(t *testing.T) {
 		t.Errorf("Production code imports the OpenTelemetry SDK (go.opentelemetry.io/otel/sdk). This is forbidden.")
 	}
 }
+
+func TestTransportTelemetry(t *testing.T) {
+	ctx := context.Background()
+	data := &TransportTelemetryData{}
+	data.SetServerAddress("localhost")
+	data.SetServerPort(8080)
+	data.SetResponseStatusCode(200)
+
+	ctx = InjectTransportTelemetry(ctx, data)
+	got := ExtractTransportTelemetry(ctx)
+	if got == nil {
+		t.Errorf("ExtractTransportTelemetry() = nil, want %v", data)
+	}
+	if got != data {
+		t.Errorf("ExtractTransportTelemetry() = %v, want %v", got, data)
+	}
+	if got.ServerAddress() != "localhost" {
+		t.Errorf("got.ServerAddress() = %q, want %q", got.ServerAddress(), "localhost")
+	}
+	if got.ServerPort() != 8080 {
+		t.Errorf("got.ServerPort() = %d, want %d", got.ServerPort(), 8080)
+	}
+	if got.ResponseStatusCode() != 200 {
+		t.Errorf("got.ResponseStatusCode() = %d, want %d", got.ResponseStatusCode(), 200)
+	}
+}
