googleapis · rahul2393 · May 8, 2025 · May 6, 2025 · May 6, 2025 · May 7, 2025
@@ -73,6 +73,11 @@ const (
 
 	// numChannels is the default value for NumChannels of client.
 	numChannels = 4
+
+	// Server timing header constants
+	serverTimingHeaderKey = "server-timing"
+	gfeTimingHeader       = "gfet4t7"
+	afeTimingHeader       = "afe"
 )
 
 const (
@@ -84,7 +89,8 @@ const (
 )
 
 var (
-	validDBPattern = regexp.MustCompile("^projects/(?P<project>[^/]+)/instances/(?P<instance>[^/]+)/databases/(?P<database>[^/]+)$")
+	validDBPattern      = regexp.MustCompile("^projects/(?P<project>[^/]+)/instances/(?P<instance>[^/]+)/databases/(?P<database>[^/]+)$")
+	serverTimingPattern = regexp.MustCompile(`([a-zA-Z0-9_-]+);\s*dur=(\d*\.?\d+)`)
 )
 
 func validDatabaseName(db string) error {
@@ -435,6 +441,28 @@ func newClientWithConfig(ctx context.Context, database string, config ClientConf
 		config.NumChannels = numChannels
 	}
 
+	var metricsProvider metric.MeterProvider
+	if emulatorAddr := os.Getenv("SPANNER_EMULATOR_HOST"); emulatorAddr != "" {
+		// Do not emit native metrics when emulator is being used
+		metricsProvider = noop.NewMeterProvider()
+	}
+	// Check if native metrics are disabled via env.
+	if disableNativeMetrics, _ := strconv.ParseBool(os.Getenv("SPANNER_DISABLE_BUILTIN_METRICS")); disableNativeMetrics {
+		config.DisableNativeMetrics = true
+	}
+	if config.DisableNativeMetrics {
+		// Do not emit native metrics when DisableNativeMetrics is set
+		metricsProvider = noop.NewMeterProvider()
+	}
+
+	metricsTracerFactory, err := newBuiltinMetricsTracerFactory(ctx, database, metricsProvider, config.Compression, opts...)
+	if err != nil {
+		return nil, err
+	}
+	if len(metricsTracerFactory.clientOpts) > 0 {
+		opts = append(opts, metricsTracerFactory.clientOpts...)
+	}
+
 	var pool gtransport.ConnPool
 
 	if gme != nil {
@@ -530,27 +558,6 @@ func newClientWithConfig(ctx context.Context, database string, config ClientConf
 	// To prevent data race in unit tests (ex: TestClient_SessionNotFound)
 	sc.mu.Lock()
 	sc.otConfig = otConfig
-	sc.mu.Unlock()
-
-	var metricsProvider metric.MeterProvider
-	if emulatorAddr := os.Getenv("SPANNER_EMULATOR_HOST"); emulatorAddr != "" {
-		// Do not emit native metrics when emulator is being used
-		metricsProvider = noop.NewMeterProvider()
-	}
-	// Check if native metrics are disabled via env.
-	if disableNativeMetrics, _ := strconv.ParseBool(os.Getenv("SPANNER_DISABLE_BUILTIN_METRICS")); disableNativeMetrics {
-		config.DisableNativeMetrics = true
-	}
-	if config.DisableNativeMetrics {
-		// Do not emit native metrics when DisableNativeMetrics is set
-		metricsProvider = noop.NewMeterProvider()
-	}
-
-	metricsTracerFactory, err := newBuiltinMetricsTracerFactory(ctx, database, metricsProvider, config.Compression, opts...)
-	if err != nil {
-		return nil, err
-	}
-	sc.mu.Lock()
 	sc.metricsTracerFactory = metricsTracerFactory
 	sc.mu.Unlock()
 
@@ -669,8 +676,9 @@ func metricsInterceptor() grpc.UnaryClientInterceptor {
 			mt.currOp.setDirectPathEnabled(true)
 		}
 
+		var md metadata.MD
 		peerInfo := &peer.Peer{}
-		opts = append(opts, grpc.Peer(peerInfo))
+		opts = append(opts, grpc.Header(&md), grpc.Peer(peerInfo))
 		err := invoker(ctx, method, req, reply, cc, opts...)
 
 		statusCode, _ := status.FromError(err)
@@ -685,6 +693,8 @@ func metricsInterceptor() grpc.UnaryClientInterceptor {
 		}
 
 		mt.currOp.currAttempt.setDirectPathUsed(isDirectPathUsed)
+		metrics := parseServerTimingHeader(md)
+		mt.currOp.currAttempt.setServerTimingMetrics(metrics)
 		recordAttemptCompletion(mt)
 		return err
 	}
@@ -1448,3 +1458,30 @@ func logf(logger *log.Logger, format string, v ...interface{}) {
 		logger.Printf(format, v...)
 	}
 }
+
+// parseServerTimingHeader extracts server timing metrics from gRPC metadata into a map
+func parseServerTimingHeader(md metadata.MD) map[string]time.Duration {
+	metrics := make(map[string]time.Duration)
+	if md == nil {
+		return metrics
+	}
+
+	serverTiming := md.Get(serverTimingHeaderKey)
+	if len(serverTiming) == 0 {
+		return metrics
+	}
+
+	for _, timing := range serverTiming {
+		matches := serverTimingPattern.FindAllStringSubmatch(timing, -1)
+		for _, match := range matches {
+			if len(match) == 3 { // full match + 2 capture groups
+				metricName := match[1]
+				duration, err := strconv.ParseFloat(match[2], 10)
+				if err == nil {
+					metrics[metricName] = time.Duration(duration*1000) * time.Microsecond
+				}
+			}
+		}
+	}
+	return metrics
+}
@@ -25,6 +25,7 @@ import (
 	"math/big"
 	"net"
 	"os"
+	"reflect"
 	"strconv"
 	"strings"
 	"sync"
@@ -45,6 +46,7 @@ import (
 	"google.golang.org/api/option"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/encoding/gzip"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/types/known/durationpb"
 	structpb "google.golang.org/protobuf/types/known/structpb"
@@ -6654,3 +6656,54 @@ func TestClient_BatchWriteExcludeTxnFromChangeStreams(t *testing.T) {
 		t.Fatal("Transaction is not set to be excluded from change streams")
 	}
 }
+
+func TestParseServerTimingHeader(t *testing.T) {
+	tests := []struct {
+		name     string
+		header   metadata.MD
+		expected map[string]time.Duration
+	}{
+		{
+			name:     "empty metadata",
+			header:   metadata.New(map[string]string{}),
+			expected: map[string]time.Duration{},
+		},
+		{
+			name:     "no server-timing header",
+			header:   metadata.New(map[string]string{"other-header": "value"}),
+			expected: map[string]time.Duration{},
+		},
+		{
+			name:     "integer duration",
+			header:   metadata.New(map[string]string{"server-timing": "gfet4t7; dur=123"}),
+			expected: map[string]time.Duration{"gfet4t7": 123 * time.Millisecond},
+		},
+		{
+			name:     "float duration",
+			header:   metadata.New(map[string]string{"server-timing": "gfet4t7; dur=123.45"}),
+			expected: map[string]time.Duration{"gfet4t7": 123*time.Millisecond + 450*time.Microsecond},
+		},
+		{
+			name:   "multiple metrics",
+			header: metadata.New(map[string]string{"server-timing": "gfet4t7; dur=123, afe; dur=456.789"}),
+			expected: map[string]time.Duration{
+				"gfet4t7": 123 * time.Millisecond,
+				"afe":     456*time.Millisecond + 789*time.Microsecond,
+			},
+		},
+		{
+			name:     "invalid duration format",
+			header:   metadata.New(map[string]string{"server-timing": "gfet4t7; dur=invalid"}),
+			expected: map[string]time.Duration{},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := parseServerTimingHeader(tt.header)
+			if !reflect.DeepEqual(got, tt.expected) {
+				t.Errorf("parseServerTimingHeader() = %v, want %v", got, tt.expected)
+			}
+		})
+	}
+}
@@ -184,7 +184,16 @@ func (g *grpcSpannerClient) ExecuteSql(ctx context.Context, req *spannerpb.Execu
 func (g *grpcSpannerClient) ExecuteStreamingSql(ctx context.Context, req *spannerpb.ExecuteSqlRequest, opts ...gax.CallOption) (spannerpb.Spanner_ExecuteStreamingSqlClient, error) {
 	// Note: This method does not add g.optsWithNextRequestID to inject x-goog-spanner-request-id
 	// as it is already manually added when creating Stream iterators for ExecuteStreamingSql.
-	return g.raw.ExecuteStreamingSql(peer.NewContext(ctx, &peer.Peer{}), req, opts...)
+	client, err := g.raw.ExecuteStreamingSql(peer.NewContext(ctx, &peer.Peer{}), req, opts...)
+	mt, ok := ctx.Value(metricsTracerKey).(*builtinMetricsTracer)
+	if !ok {
+		return client, err
+	}
+	if mt != nil && client != nil && mt.currOp.currAttempt != nil {
+		md, _ := client.Header()
+		mt.currOp.currAttempt.setServerTimingMetrics(parseServerTimingHeader(md))
+	}
+	return client, err
 }
 
 func (g *grpcSpannerClient) ExecuteBatchDml(ctx context.Context, req *spannerpb.ExecuteBatchDmlRequest, opts ...gax.CallOption) (*spannerpb.ExecuteBatchDmlResponse, error) {
@@ -210,7 +219,16 @@ func (g *grpcSpannerClient) Read(ctx context.Context, req *spannerpb.ReadRequest
 func (g *grpcSpannerClient) StreamingRead(ctx context.Context, req *spannerpb.ReadRequest, opts ...gax.CallOption) (spannerpb.Spanner_StreamingReadClient, error) {
 	// Note: This method does not add g.optsWithNextRequestID, as it is already
 	// manually added when creating Stream iterators for StreamingRead.
-	return g.raw.StreamingRead(peer.NewContext(ctx, &peer.Peer{}), req, opts...)
+	client, err := g.raw.StreamingRead(peer.NewContext(ctx, &peer.Peer{}), req, opts...)
+	mt, ok := ctx.Value(metricsTracerKey).(*builtinMetricsTracer)
+	if !ok {
+		return client, err
+	}
+	if mt != nil && client != nil && mt.currOp.currAttempt != nil {
+		md, _ := client.Header()
+		mt.currOp.currAttempt.setServerTimingMetrics(parseServerTimingHeader(md))
+	}
+	return client, err
 }
 
 func (g *grpcSpannerClient) BeginTransaction(ctx context.Context, req *spannerpb.BeginTransactionRequest, opts ...gax.CallOption) (*spannerpb.Transaction, error) {
@@ -264,5 +282,14 @@ func (g *grpcSpannerClient) PartitionRead(ctx context.Context, req *spannerpb.Pa
 }
 
 func (g *grpcSpannerClient) BatchWrite(ctx context.Context, req *spannerpb.BatchWriteRequest, opts ...gax.CallOption) (spannerpb.Spanner_BatchWriteClient, error) {
-	return g.raw.BatchWrite(peer.NewContext(ctx, &peer.Peer{}), req, g.optsWithNextRequestID(opts)...)
+	client, err := g.raw.BatchWrite(peer.NewContext(ctx, &peer.Peer{}), req, g.optsWithNextRequestID(opts)...)
+	mt, ok := ctx.Value(metricsTracerKey).(*builtinMetricsTracer)
+	if !ok {
+		return client, err
+	}
+	if mt != nil && client != nil && mt.currOp.currAttempt != nil {
+		md, _ := client.Header()
+		mt.currOp.currAttempt.setServerTimingMetrics(parseServerTimingHeader(md))
+	}
+	return client, err
 }
@@ -1087,6 +1087,10 @@ func (s *inMemSpannerServer) StreamingRead(req *spannerpb.ReadRequest, stream sp
 			req.Table,
 		),
 	}
+	header := metadata.New(map[string]string{"server-timing": "gfet4t7; dur=123"})
+	if err := grpc.SendHeader(stream.Context(), header); err != nil {
+		return gstatus.Errorf(codes.Internal, "unable to send 'server-timing' header")
+	}
 	return s.executeStreamingSQL(sqlReq, stream)
 }