client: restore remote address in traces (#2718)

The client-side traces were otherwise only showing `RPC: to <nil>`,
which is not helpful.

Also clean up construction of traceInfo and firstLine in a few places.

internal/transport/http2_client.go

@@ -1369,6 +1369,8 @@ func (t *http2Client) ChannelzMetric() *channelz.SocketInternalMetric {
 	return &s
 }
 
+func (t *http2Client) RemoteAddr() net.Addr { return t.remoteAddr }
+
 func (t *http2Client) IncrMsgSent() {
 	atomic.AddInt64(&t.czData.msgSent, 1)
 	atomic.StoreInt64(&t.czData.lastMsgSentTime, time.Now().UnixNano())

internal/transport/transport.go

@@ -610,6 +610,9 @@ type ClientTransport interface {
 	// GetGoAwayReason returns the reason why GoAway frame was received.
 	GetGoAwayReason() GoAwayReason
 
+	// RemoteAddr returns the remote network address.
+	RemoteAddr() net.Addr
+
 	// IncrMsgSent increments the number of message sent through this transport.
 	IncrMsgSent()
 

server.go

@@ -749,10 +749,11 @@ func (s *Server) traceInfo(st transport.ServerTransport, stream *transport.Strea
 
 	trInfo = &traceInfo{
 		tr: tr,
+		firstLine: firstLine{
+			client:     false,
+			remoteAddr: st.RemoteAddr(),
+		},
 	}
-	trInfo.firstLine.client = false
-	trInfo.firstLine.remoteAddr = st.RemoteAddr()
-
 	if dl, ok := stream.Context().Deadline(); ok {
 		trInfo.firstLine.deadline = time.Until(dl)
 	}
@@ -860,7 +861,6 @@ func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.
 	}
 	if trInfo != nil {
 		defer trInfo.tr.Finish()
-		trInfo.firstLine.client = false
 		trInfo.tr.LazyLog(&trInfo.firstLine, false)
 		defer func() {
 			if err != nil && err != io.EOF {

stream.go

@@ -230,10 +230,14 @@ func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, meth
 	if c.creds != nil {
 		callHdr.Creds = c.creds
 	}
-	var trInfo traceInfo
+	var trInfo *traceInfo
 	if EnableTracing {
-		trInfo.tr = trace.New("grpc.Sent."+methodFamily(method), method)
-		trInfo.firstLine.client = true
+		trInfo = &traceInfo{
+			tr: trace.New("grpc.Sent."+methodFamily(method), method),
+			firstLine: firstLine{
+				client: true,
+			},
+		}
 		if deadline, ok := ctx.Deadline(); ok {
 			trInfo.firstLine.deadline = time.Until(deadline)
 		}
@@ -323,7 +327,7 @@ func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, meth
 	return cs, nil
 }
 
-func (cs *clientStream) newAttemptLocked(sh stats.Handler, trInfo traceInfo) error {
+func (cs *clientStream) newAttemptLocked(sh stats.Handler, trInfo *traceInfo) error {
 	cs.attempt = &csAttempt{
 		cs:           cs,
 		dc:           cs.cc.dopts.dc,
@@ -338,6 +342,9 @@ func (cs *clientStream) newAttemptLocked(sh stats.Handler, trInfo traceInfo) err
 	if err != nil {
 		return err
 	}
+	if trInfo != nil {
+		trInfo.firstLine.SetRemoteAddr(t.RemoteAddr())
+	}
 	cs.attempt.t = t
 	cs.attempt.done = done
 	return nil
@@ -414,9 +421,10 @@ type csAttempt struct {
 	decompSet bool
 
 	mu sync.Mutex // guards trInfo.tr
+	// trInfo may be nil (if EnableTracing is false).
 	// trInfo.tr is set when created (if EnableTracing is true),
 	// and cleared when the finish method is called.
-	trInfo traceInfo
+	trInfo *traceInfo
 
 	statsHandler stats.Handler
 }
@@ -540,7 +548,7 @@ func (cs *clientStream) retryLocked(lastErr error) error {
 			cs.commitAttemptLocked()
 			return err
 		}
-		if err := cs.newAttemptLocked(nil, traceInfo{}); err != nil {
+		if err := cs.newAttemptLocked(nil, nil); err != nil {
 			return err
 		}
 		if lastErr = cs.replayBufferLocked(); lastErr == nil {
@@ -811,7 +819,7 @@ func (cs *clientStream) finish(err error) {
 
 func (a *csAttempt) sendMsg(m interface{}, hdr, payld, data []byte) error {
 	cs := a.cs
-	if EnableTracing {
+	if a.trInfo != nil {
 		a.mu.Lock()
 		if a.trInfo.tr != nil {
 			a.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
@@ -868,7 +876,7 @@ func (a *csAttempt) recvMsg(m interface{}, payInfo *payloadInfo) (err error) {
 		}
 		return toRPCErr(err)
 	}
-	if EnableTracing {
+	if a.trInfo != nil {
 		a.mu.Lock()
 		if a.trInfo.tr != nil {
 			a.trInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
@@ -944,7 +952,7 @@ func (a *csAttempt) finish(err error) {
 		}
 		a.statsHandler.HandleRPC(a.cs.ctx, end)
 	}
-	if a.trInfo.tr != nil {
+	if a.trInfo != nil && a.trInfo.tr != nil {
 		if err == nil {
 			a.trInfo.tr.LazyPrintf("RPC: [OK]")
 		} else {

trace.go

@@ -24,6 +24,7 @@ import (
 	"io"
 	"net"
 	"strings"
+	"sync"
 	"time"
 
 	"golang.org/x/net/trace"
@@ -53,13 +54,25 @@ type traceInfo struct {
 }
 
 // firstLine is the first line of an RPC trace.
+// It may be mutated after construction; remoteAddr specifically may change
+// during client-side use.
 type firstLine struct {
+	mu         sync.Mutex
 	client     bool // whether this is a client (outgoing) RPC
 	remoteAddr net.Addr
 	deadline   time.Duration // may be zero
 }
 
+func (f *firstLine) SetRemoteAddr(addr net.Addr) {
+	f.mu.Lock()
+	f.remoteAddr = addr
+	f.mu.Unlock()
+}
+
 func (f *firstLine) String() string {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
 	var line bytes.Buffer
 	io.WriteString(&line, "RPC: ")
 	if f.client {