diff --git a/v2/telemetry.go b/v2/telemetry.go index ee990838..7424057c 100644 --- 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 index 28741898..fa2c7514 100644 --- 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) + } +}