From 0252e09f2156f21768e35b23d760fc46eb649d15 Mon Sep 17 00:00:00 2001
From: yangzhouhan <yangzhouhan@gmail.com>
Date: Thu, 23 Jul 2015 10:55:42 -0700
Subject: [PATCH] add response trace log for client streaming

---
 call.go   | 4 ++--
 stream.go | 9 ++++++++-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/call.go b/call.go
index 3b974721..218fc0f6 100644
--- a/call.go
+++ b/call.go
@@ -165,7 +165,7 @@ func Invoke(ctx context.Context, method string, args, reply interface{}, cc *Cli
 			return toRPCErr(err)
 		}
 		if EnableTracing {
-			c.traceInfo.tr.LazyLog(payload{args}, true)
+			c.traceInfo.tr.LazyLog(&fmtStringer{"sent: %v", []interface{}{payload{args}}}, true)
 		}
 		stream, err = sendRequest(ctx, cc.dopts.codec, callHdr, t, args, topts)
 		if err != nil {
@@ -184,7 +184,7 @@ func Invoke(ctx context.Context, method string, args, reply interface{}, cc *Cli
 			continue
 		}
 		if EnableTracing {
-			c.traceInfo.tr.LazyLog(payload{reply}, true)
+			c.traceInfo.tr.LazyLog(&fmtStringer{"received: %v", []interface{}{payload{reply}}}, true)
 		}
 		t.CloseStream(stream, lastErr)
 		if lastErr != nil {
diff --git a/stream.go b/stream.go
index c902034d..348fe9a5 100644
--- a/stream.go
+++ b/stream.go
@@ -166,7 +166,7 @@ func (cs *clientStream) SendMsg(m interface{}) (err error) {
 	if cs.tracing {
 		cs.mu.Lock()
 		if cs.traceInfo.tr != nil {
-			cs.traceInfo.tr.LazyLog(payload{m}, true)
+			cs.traceInfo.tr.LazyLog(&fmtStringer{"sent: %v", []interface{}{payload{m}}}, true)
 		}
 		cs.mu.Unlock()
 	}
@@ -187,6 +187,13 @@ func (cs *clientStream) SendMsg(m interface{}) (err error) {
 }
 
 func (cs *clientStream) RecvMsg(m interface{}) (err error) {
+	if cs.tracing {
+		cs.mu.Lock()
+		if cs.traceInfo.tr != nil {
+			cs.traceInfo.tr.LazyLog(&fmtStringer{"received: %v", []interface{}{payload{m}}}, true)
+		}
+		cs.mu.Unlock()
+	}
 	err = recv(cs.p, cs.codec, m)
 	defer func() {
 		// err != nil indicates the termination of the stream.