fix wrong context when handling stats

call_test.go

@@ -185,6 +185,8 @@ func (s *server) start(t *testing.T, port int, maxStreams uint32) {
 		}
 		go st.HandleStreams(func(s *transport.Stream) {
 			go h.handleStream(t, s)
+		}, func(ctx context.Context, method string) context.Context {
+			return ctx
 		})
 	}
 }

server.go

@@ -467,6 +467,12 @@ func (s *Server) serveStreams(st transport.ServerTransport) {
 			defer wg.Done()
 			s.handleStream(st, stream, s.traceInfo(st, stream))
 		}()
+	}, func(ctx context.Context, method string) context.Context {
+		if !EnableTracing {
+			return ctx
+		}
+		tr := trace.New("grpc.Recv."+methodFamily(method), method)
+		return trace.NewContext(ctx, tr)
 	})
 	wg.Wait()
 }
@@ -519,12 +525,17 @@ func (s *Server) traceInfo(st transport.ServerTransport, stream *transport.Strea
 	if !EnableTracing {
 		return nil
 	}
+	tr, ok := trace.FromContext(stream.Context())
+	if !ok {
+		grpclog.Fatalf("cannot get trace from context while EnableTracing == true")
+	}
+
 	trInfo = &traceInfo{
-		tr: trace.New("grpc.Recv."+methodFamily(stream.Method()), stream.Method()),
+		tr: tr,
 	}
 	trInfo.firstLine.client = false
 	trInfo.firstLine.remoteAddr = st.RemoteAddr()
-	stream.TraceContext(trInfo.tr)
+
 	if dl, ok := stream.Context().Deadline(); ok {
 		trInfo.firstLine.deadline = dl.Sub(time.Now())
 	}

stats/stats_test.go

@@ -592,6 +592,12 @@ func TestServerStatsUnaryRPC(t *testing.T) {
 		t.Fatalf("got %v stats, want %v stats", len(got), len(checkFuncs))
 	}
 
+	for i := 0; i < len(got)-1; i++ {
+		if got[i].ctx != got[i+1].ctx {
+			t.Fatalf("got different contexts with two stats %T %T", got[i].s, got[i+1].s)
+		}
+	}
+
 	for i, f := range checkFuncs {
 		mu.Lock()
 		f(t, got[i], expect)
@@ -645,6 +651,12 @@ func TestServerStatsUnaryRPCError(t *testing.T) {
 		t.Fatalf("got %v stats, want %v stats", len(got), len(checkFuncs))
 	}
 
+	for i := 0; i < len(got)-1; i++ {
+		if got[i].ctx != got[i+1].ctx {
+			t.Fatalf("got different contexts with two stats %T %T", got[i].s, got[i+1].s)
+		}
+	}
+
 	for i, f := range checkFuncs {
 		mu.Lock()
 		f(t, got[i], expect)
@@ -704,6 +716,12 @@ func TestServerStatsStreamingRPC(t *testing.T) {
 		t.Fatalf("got %v stats, want %v stats", len(got), len(checkFuncs))
 	}
 
+	for i := 0; i < len(got)-1; i++ {
+		if got[i].ctx != got[i+1].ctx {
+			t.Fatalf("got different contexts with two stats %T %T", got[i].s, got[i+1].s)
+		}
+	}
+
 	for i, f := range checkFuncs {
 		mu.Lock()
 		f(t, got[i], expect)
@@ -759,6 +777,12 @@ func TestServerStatsStreamingRPCError(t *testing.T) {
 		t.Fatalf("got %v stats, want %v stats", len(got), len(checkFuncs))
 	}
 
+	for i := 0; i < len(got)-1; i++ {
+		if got[i].ctx != got[i+1].ctx {
+			t.Fatalf("got different contexts with two stats %T %T", got[i].s, got[i+1].s)
+		}
+	}
+
 	for i, f := range checkFuncs {
 		mu.Lock()
 		f(t, got[i], expect)
@@ -923,6 +947,12 @@ func TestClientStatsUnaryRPCError(t *testing.T) {
 		t.Fatalf("got %v stats, want %v stats", len(got), len(checkFuncs))
 	}
 
+	for i := 0; i < len(got)-1; i++ {
+		if got[i].ctx != got[i+1].ctx {
+			t.Fatalf("got different contexts with two stats %T %T", got[i].s, got[i+1].s)
+		}
+	}
+
 	for i, f := range checkFuncs {
 		mu.Lock()
 		f(t, got[i], expect)

transport/handler_server.go

@@ -268,7 +268,7 @@ func (ht *serverHandlerTransport) WriteHeader(s *Stream, md metadata.MD) error {
 	})
 }
 
-func (ht *serverHandlerTransport) HandleStreams(startStream func(*Stream)) {
+func (ht *serverHandlerTransport) HandleStreams(startStream func(*Stream), traceCtx func(context.Context, string) context.Context) {
 	// With this transport type there will be exactly 1 stream: this HTTP request.
 
 	var ctx context.Context

transport/http2_server.go

@@ -151,7 +151,7 @@ func newHTTP2Server(conn net.Conn, config *ServerConfig) (_ ServerTransport, err
 }
 
 // operateHeader takes action on the decoded headers.
-func (t *http2Server) operateHeaders(frame *http2.MetaHeadersFrame, handle func(*Stream)) (close bool) {
+func (t *http2Server) operateHeaders(frame *http2.MetaHeadersFrame, handle func(*Stream), traceCtx func(context.Context, string) context.Context) (close bool) {
 	buf := newRecvBuffer()
 	s := &Stream{
 		id:  frame.Header().StreamID,
@@ -239,6 +239,7 @@ func (t *http2Server) operateHeaders(frame *http2.MetaHeadersFrame, handle func(
 	s.windowHandler = func(n int) {
 		t.updateWindow(s, uint32(n))
 	}
+	s.ctx = traceCtx(s.ctx, s.method)
 	if stats.On() {
 		inHeader := &stats.InHeader{
 			FullMethod: s.method,
@@ -255,7 +256,8 @@ func (t *http2Server) operateHeaders(frame *http2.MetaHeadersFrame, handle func(
 
 // HandleStreams receives incoming streams using the given handler. This is
 // typically run in a separate goroutine.
-func (t *http2Server) HandleStreams(handle func(*Stream)) {
+// traceCtx attaches trace to ctx and returns the new context.
+func (t *http2Server) HandleStreams(handle func(*Stream), traceCtx func(context.Context, string) context.Context) {
 	// Check the validity of client preface.
 	preface := make([]byte, len(clientPreface))
 	if _, err := io.ReadFull(t.conn, preface); err != nil {
@@ -310,7 +312,7 @@ func (t *http2Server) HandleStreams(handle func(*Stream)) {
 		}
 		switch frame := frame.(type) {
 		case *http2.MetaHeadersFrame:
-			if t.operateHeaders(frame, handle) {
+			if t.operateHeaders(frame, handle, traceCtx) {
 				t.Close()
 				break
 			}

transport/transport.go

@@ -479,7 +479,7 @@ type ClientTransport interface {
 // Write methods for a given Stream will be called serially.
 type ServerTransport interface {
 	// HandleStreams receives incoming streams using the given handler.
-	HandleStreams(func(*Stream))
+	HandleStreams(func(*Stream), func(context.Context, string) context.Context)
 
 	// WriteHeader sends the header metadata for the given stream.
 	// WriteHeader may not be called on all streams.