From 2c2d834e8e13f239343a3441b76b0efcc167a0d7 Mon Sep 17 00:00:00 2001
From: Brian Tiger Chow <734339+btc@users.noreply.github.com>
Date: Mon, 12 Mar 2018 16:16:36 +0000
Subject: [PATCH] stats: add BeginTime to stats.End (#1907)

---
 server.go           | 12 ++++++++----
 stats/stats.go      |  2 ++
 stats/stats_test.go |  3 +++
 stream.go           | 19 +++++++++++++------
 4 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/server.go b/server.go
index c085b23c..997b24fa 100644
--- a/server.go
+++ b/server.go
@@ -777,13 +777,15 @@ func (s *Server) sendResponse(t transport.ServerTransport, stream *transport.Str
 func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.Stream, srv *service, md *MethodDesc, trInfo *traceInfo) (err error) {
 	sh := s.opts.statsHandler
 	if sh != nil {
+		beginTime := time.Now()
 		begin := &stats.Begin{
-			BeginTime: time.Now(),
+			BeginTime: beginTime,
 		}
 		sh.HandleRPC(stream.Context(), begin)
 		defer func() {
 			end := &stats.End{
-				EndTime: time.Now(),
+				BeginTime: beginTime,
+				EndTime:   time.Now(),
 			}
 			if err != nil && err != io.EOF {
 				end.Error = toRPCErr(err)
@@ -977,13 +979,15 @@ func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.
 func (s *Server) processStreamingRPC(t transport.ServerTransport, stream *transport.Stream, srv *service, sd *StreamDesc, trInfo *traceInfo) (err error) {
 	sh := s.opts.statsHandler
 	if sh != nil {
+		beginTime := time.Now()
 		begin := &stats.Begin{
-			BeginTime: time.Now(),
+			BeginTime: beginTime,
 		}
 		sh.HandleRPC(stream.Context(), begin)
 		defer func() {
 			end := &stats.End{
-				EndTime: time.Now(),
+				BeginTime: beginTime,
+				EndTime:   time.Now(),
 			}
 			if err != nil && err != io.EOF {
 				end.Error = toRPCErr(err)
diff --git a/stats/stats.go b/stats/stats.go
index d5aa2f79..3f13190a 100644
--- a/stats/stats.go
+++ b/stats/stats.go
@@ -169,6 +169,8 @@ func (s *OutTrailer) isRPCStats() {}
 type End struct {
 	// Client is true if this End is from client side.
 	Client bool
+	// BeginTime is the time when the RPC began.
+	BeginTime time.Time
 	// EndTime is the time when the RPC ends.
 	EndTime time.Time
 	// Error is the error the RPC ended with. It is an error generated from
diff --git a/stats/stats_test.go b/stats/stats_test.go
index b6c7b998..00a5e4f3 100644
--- a/stats/stats_test.go
+++ b/stats/stats_test.go
@@ -638,6 +638,9 @@ func checkEnd(t *testing.T, d *gotData, e *expectedData) {
 	if d.ctx == nil {
 		t.Fatalf("d.ctx = nil, want <non-nil>")
 	}
+	if st.BeginTime.IsZero() {
+		t.Fatalf("st.BeginTime = %v, want <non-zero>", st.BeginTime)
+	}
 	if st.EndTime.IsZero() {
 		t.Fatalf("st.EndTime = %v, want <non-zero>", st.EndTime)
 	}
diff --git a/stream.go b/stream.go
index c5cb9bc7..64e66e17 100644
--- a/stream.go
+++ b/stream.go
@@ -205,11 +205,13 @@ func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, meth
 	}
 	ctx = newContextWithRPCInfo(ctx, c.failFast)
 	sh := cc.dopts.copts.StatsHandler
+	var beginTime time.Time
 	if sh != nil {
 		ctx = sh.TagRPC(ctx, &stats.RPCTagInfo{FullMethodName: method, FailFast: c.failFast})
+		beginTime = time.Now()
 		begin := &stats.Begin{
 			Client:    true,
-			BeginTime: time.Now(),
+			BeginTime: beginTime,
 			FailFast:  c.failFast,
 		}
 		sh.HandleRPC(ctx, begin)
@@ -217,8 +219,10 @@ func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, meth
 			if err != nil {
 				// Only handle end stats if err != nil.
 				end := &stats.End{
-					Client: true,
-					Error:  err,
+					Client:    true,
+					Error:     err,
+					BeginTime: beginTime,
+					EndTime:   time.Now(),
 				}
 				sh.HandleRPC(ctx, end)
 			}
@@ -283,6 +287,7 @@ func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, meth
 
 		statsCtx:     ctx,
 		statsHandler: cc.dopts.copts.StatsHandler,
+		beginTime:    beginTime,
 	}
 	if desc != unaryStreamDesc {
 		// Listen on cc and stream contexts to cleanup when the user closes the
@@ -337,6 +342,7 @@ type clientStream struct {
 	// so that all the generated stats for a particular RPC can be associated in the processing phase.
 	statsCtx     context.Context
 	statsHandler stats.Handler
+	beginTime    time.Time
 }
 
 func (cs *clientStream) Context() context.Context {
@@ -512,9 +518,10 @@ func (cs *clientStream) finish(err error) {
 	}
 	if cs.statsHandler != nil {
 		end := &stats.End{
-			Client:  true,
-			EndTime: time.Now(),
-			Error:   err,
+			Client:    true,
+			BeginTime: cs.beginTime,
+			EndTime:   time.Now(),
+			Error:     err,
 		}
 		cs.statsHandler.HandleRPC(cs.statsCtx, end)
 	}