wrap stream for tracing

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>

grpc.go

@@ -51,6 +51,18 @@ type ServerReflection struct {
 }
 */
 
+type streamWrapper struct {
+	ctx context.Context
+	grpc.ServerStream
+}
+
+func (w *streamWrapper) Context() context.Context {
+	if w.ctx != nil {
+		return w.ctx
+	}
+	return w.ServerStream.Context()
+}
+
 type Server struct {
 	handlers       map[string]server.Handler
 	srv            *grpc.Server
@@ -206,6 +218,8 @@ func (g *Server) handler(srv interface{}, stream grpc.ServerStream) error {
 		),
 	)
 
+	stream = &streamWrapper{ctx, stream}
+
 	ts := time.Now()
 	g.opts.Meter.Counter(semconv.ServerRequestInflight, "endpoint", fullMethod).Inc()