diff options
Diffstat (limited to 'vendor/google.golang.org/grpc/call.go')
| -rw-r--r-- | vendor/google.golang.org/grpc/call.go | 29 |
1 files changed, 17 insertions, 12 deletions
diff --git a/vendor/google.golang.org/grpc/call.go b/vendor/google.golang.org/grpc/call.go index 4d8023d..81b52be 100644 --- a/vendor/google.golang.org/grpc/call.go +++ b/vendor/google.golang.org/grpc/call.go @@ -42,6 +42,7 @@ import ( "golang.org/x/net/context" "golang.org/x/net/trace" "google.golang.org/grpc/codes" + "google.golang.org/grpc/peer" "google.golang.org/grpc/stats" "google.golang.org/grpc/transport" ) @@ -66,7 +67,7 @@ func recvResponse(ctx context.Context, dopts dialOptions, t transport.ClientTran } p := &parser{r: stream} var inPayload *stats.InPayload - if stats.On() { + if dopts.copts.StatsHandler != nil { inPayload = &stats.InPayload{ Client: true, } @@ -82,14 +83,17 @@ func recvResponse(ctx context.Context, dopts dialOptions, t transport.ClientTran if inPayload != nil && err == io.EOF && stream.StatusCode() == codes.OK { // TODO in the current implementation, inTrailer may be handled before inPayload in some cases. // Fix the order if necessary. - stats.HandleRPC(ctx, inPayload) + dopts.copts.StatsHandler.HandleRPC(ctx, inPayload) } c.trailerMD = stream.Trailer() + if peer, ok := peer.FromContext(stream.Context()); ok { + c.peer = peer + } return nil } // sendRequest writes out various information of an RPC such as Context and Message. -func sendRequest(ctx context.Context, codec Codec, compressor Compressor, callHdr *transport.CallHdr, t transport.ClientTransport, args interface{}, opts *transport.Options) (_ *transport.Stream, err error) { +func sendRequest(ctx context.Context, dopts dialOptions, compressor Compressor, callHdr *transport.CallHdr, t transport.ClientTransport, args interface{}, opts *transport.Options) (_ *transport.Stream, err error) { stream, err := t.NewStream(ctx, callHdr) if err != nil { return nil, err @@ -109,19 +113,19 @@ func sendRequest(ctx context.Context, codec Codec, compressor Compressor, callHd if compressor != nil { cbuf = new(bytes.Buffer) } - if stats.On() { + if dopts.copts.StatsHandler != nil { outPayload = &stats.OutPayload{ Client: true, } } - outBuf, err := encode(codec, args, compressor, cbuf, outPayload) + outBuf, err := encode(dopts.codec, args, compressor, cbuf, outPayload) if err != nil { return nil, Errorf(codes.Internal, "grpc: %v", err) } err = t.Write(stream, outBuf, opts) if err == nil && outPayload != nil { outPayload.SentTime = time.Now() - stats.HandleRPC(ctx, outPayload) + dopts.copts.StatsHandler.HandleRPC(ctx, outPayload) } // t.NewStream(...) could lead to an early rejection of the RPC (e.g., the service/method // does not exist.) so that t.Write could get io.EOF from wait(...). Leave the following @@ -179,23 +183,24 @@ func invoke(ctx context.Context, method string, args, reply interface{}, cc *Cli } }() } - if stats.On() { - ctx = stats.TagRPC(ctx, &stats.RPCTagInfo{FullMethodName: method}) + sh := cc.dopts.copts.StatsHandler + if sh != nil { + ctx = sh.TagRPC(ctx, &stats.RPCTagInfo{FullMethodName: method}) begin := &stats.Begin{ Client: true, BeginTime: time.Now(), FailFast: c.failFast, } - stats.HandleRPC(ctx, begin) + sh.HandleRPC(ctx, begin) } defer func() { - if stats.On() { + if sh != nil { end := &stats.End{ Client: true, EndTime: time.Now(), Error: e, } - stats.HandleRPC(ctx, end) + sh.HandleRPC(ctx, end) } }() topts := &transport.Options{ @@ -241,7 +246,7 @@ func invoke(ctx context.Context, method string, args, reply interface{}, cc *Cli if c.traceInfo.tr != nil { c.traceInfo.tr.LazyLog(&payload{sent: true, msg: args}, true) } - stream, err = sendRequest(ctx, cc.dopts.codec, cc.dopts.cp, callHdr, t, args, topts) + stream, err = sendRequest(ctx, cc.dopts, cc.dopts.cp, callHdr, t, args, topts) if err != nil { if put != nil { put() |