Skip to content

Commit 88a36cd

Browse files
rahul2393rahul2393
authored
fix(spanner): remove stream wrapper for direct path check (#12622)
* fix(spanner): remove stream wrapper for direct path check * refactoring
1 parent e455b54 commit 88a36cd

File tree

3 files changed

+21
-60
lines changed

3 files changed

+21
-60
lines changed

‎spanner/client.go‎

Lines changed: 9 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ package spanner
1818

1919
import (
2020
"context"
21-
"errors"
2221
"fmt"
2322
"io"
2423
"log"
@@ -718,69 +717,14 @@ func metricsInterceptor() grpc.UnaryClientInterceptor {
718717

719718
statusCode, _ := status.FromError(err)
720719
mt.currOp.currAttempt.setStatus(statusCode.Code().String())
721-
722-
isDirectPathUsed := false
723-
if peerInfo.Addr != nil {
724-
remoteIP := peerInfo.Addr.String()
725-
if strings.HasPrefix(remoteIP, directPathIPV4Prefix) || strings.HasPrefix(remoteIP, directPathIPV6Prefix) {
726-
isDirectPathUsed = true
727-
}
728-
}
729-
730-
mt.currOp.currAttempt.setDirectPathUsed(isDirectPathUsed)
720+
mt.currOp.currAttempt.setDirectPathUsed(peer.NewContext(ctx, peerInfo))
731721
metrics := parseServerTimingHeader(md)
732722
mt.currOp.currAttempt.setServerTimingMetrics(metrics)
733723
recordAttemptCompletion(mt)
734724
return err
735725
}
736726
}
737727

738-
// wrappedStream wraps around the embedded grpc.ClientStream, and intercepts the RecvMsg and
739-
// SendMsg method call.
740-
type wrappedStream struct {
741-
method string
742-
target string
743-
grpc.ClientStream
744-
}
745-
746-
func (w *wrappedStream) RecvMsg(m any) error {
747-
err := w.ClientStream.RecvMsg(m)
748-
if errors.Is(err, io.EOF) {
749-
return err
750-
}
751-
ctx := w.ClientStream.Context()
752-
mt, ok := ctx.Value(metricsTracerKey).(*builtinMetricsTracer)
753-
if !ok {
754-
return err
755-
}
756-
mt.method = w.method
757-
if strings.HasPrefix(w.target, "google-c2p") {
758-
mt.currOp.setDirectPathEnabled(true)
759-
}
760-
isDirectPathUsed := false
761-
peerInfo, ok := peer.FromContext(ctx)
762-
if ok {
763-
if peerInfo.Addr != nil {
764-
remoteIP := peerInfo.Addr.String()
765-
if strings.HasPrefix(remoteIP, directPathIPV4Prefix) || strings.HasPrefix(remoteIP, directPathIPV6Prefix) {
766-
isDirectPathUsed = true
767-
}
768-
}
769-
}
770-
if mt.currOp.currAttempt != nil {
771-
mt.currOp.currAttempt.setDirectPathUsed(isDirectPathUsed)
772-
}
773-
return err
774-
}
775-
776-
func (w *wrappedStream) SendMsg(m any) error {
777-
return w.ClientStream.SendMsg(m)
778-
}
779-
780-
func newWrappedStream(s grpc.ClientStream, method, target string) grpc.ClientStream {
781-
return &wrappedStream{ClientStream: s, method: method, target: target}
782-
}
783-
784728
// metricsInterceptor is a gRPC stream client interceptor that records metrics for stream RPCs.
785729
func metricsStreamInterceptor() grpc.StreamClientInterceptor {
786730
return func(
@@ -795,7 +739,14 @@ func metricsStreamInterceptor() grpc.StreamClientInterceptor {
795739
if err != nil {
796740
return nil, err
797741
}
798-
return newWrappedStream(s, method, cc.Target()), nil
742+
mt, ok := ctx.Value(metricsTracerKey).(*builtinMetricsTracer)
743+
if ok && mt != nil {
744+
mt.method = method
745+
if strings.HasPrefix(cc.Target(), "google-c2p") {
746+
mt.currOp.setDirectPathEnabled(true)
747+
}
748+
}
749+
return s, nil
799750
}
800751
}
801752

‎spanner/grpc_client.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,7 @@ func (g *grpcSpannerClient) ExecuteStreamingSql(ctx context.Context, req *spanne
192192
if mt != nil && client != nil && mt.currOp.currAttempt != nil {
193193
md, _ := client.Header()
194194
mt.currOp.currAttempt.setServerTimingMetrics(parseServerTimingHeader(md))
195+
mt.currOp.currAttempt.setDirectPathUsed(client.Context())
195196
}
196197
return client, err
197198
}
@@ -227,6 +228,7 @@ func (g *grpcSpannerClient) StreamingRead(ctx context.Context, req *spannerpb.Re
227228
if mt != nil && client != nil && mt.currOp.currAttempt != nil {
228229
md, _ := client.Header()
229230
mt.currOp.currAttempt.setServerTimingMetrics(parseServerTimingHeader(md))
231+
mt.currOp.currAttempt.setDirectPathUsed(client.Context())
230232
}
231233
return client, err
232234
}
@@ -290,6 +292,7 @@ func (g *grpcSpannerClient) BatchWrite(ctx context.Context, req *spannerpb.Batch
290292
if mt != nil && client != nil && mt.currOp.currAttempt != nil {
291293
md, _ := client.Header()
292294
mt.currOp.currAttempt.setServerTimingMetrics(parseServerTimingHeader(md))
295+
mt.currOp.currAttempt.setDirectPathUsed(client.Context())
293296
}
294297
return client, err
295298
}

‎spanner/metrics.go‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ import (
3838
"google.golang.org/grpc"
3939
"google.golang.org/grpc/codes"
4040
"google.golang.org/grpc/experimental/stats"
41+
"google.golang.org/grpc/peer"
4142
"google.golang.org/grpc/stats/opentelemetry"
4243
"google.golang.org/grpc/status"
4344

@@ -508,8 +509,14 @@ func (o *opTracer) incrementAttemptCount() {
508509
}
509510

510511
// setDirectPathUsed sets whether DirectPath was used for the attempt.
511-
func (a *attemptTracer) setDirectPathUsed(used bool) {
512-
a.directPathUsed = used
512+
func (a *attemptTracer) setDirectPathUsed(ctx context.Context) {
513+
peerInfo, ok := peer.FromContext(ctx)
514+
if ok && peerInfo.Addr != nil {
515+
remoteIP := peerInfo.Addr.String()
516+
if strings.HasPrefix(remoteIP, directPathIPV4Prefix) || strings.HasPrefix(remoteIP, directPathIPV6Prefix) {
517+
a.directPathUsed = true
518+
}
519+
}
513520
}
514521

515522
func (a *attemptTracer) setServerTimingMetrics(metrics map[string]time.Duration) {

0 commit comments

Comments
 (0)