[#15] tracing: Add events for grpc stream calls

Signed-off-by: Roman Loginov <r.loginov@yadro.com>
This commit is contained in:
Roman Loginov 2024-11-20 11:35:34 +03:00
parent f17779933e
commit 37bd758211
3 changed files with 44 additions and 6 deletions

View file

@ -2,7 +2,10 @@ package grpc
import (
"context"
"fmt"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
)
@ -10,14 +13,16 @@ import (
type clientStream struct {
originalStream grpc.ClientStream
desc *grpc.StreamDesc
span trace.Span
finished chan<- error
done <-chan struct{}
}
func newgRPCClientStream(originalStream grpc.ClientStream, desc *grpc.StreamDesc, finished chan<- error, done <-chan struct{}) grpc.ClientStream {
func newgRPCClientStream(originalStream grpc.ClientStream, desc *grpc.StreamDesc, span trace.Span, finished chan<- error, done <-chan struct{}) grpc.ClientStream {
return &clientStream{
originalStream: originalStream,
desc: desc,
span: span,
finished: finished,
done: done,
}
@ -39,6 +44,7 @@ func (cs *clientStream) Trailer() metadata.MD {
}
func (cs *clientStream) CloseSend() error {
cs.span.AddEvent("client.stream.close.send.start")
err := cs.originalStream.CloseSend()
if err != nil {
select {
@ -46,6 +52,7 @@ func (cs *clientStream) CloseSend() error {
case cs.finished <- err:
}
}
cs.span.AddEvent("client.stream.close.send.finish")
return err
}
@ -54,6 +61,9 @@ func (cs *clientStream) Context() context.Context {
}
func (cs *clientStream) SendMsg(m any) error {
cs.span.AddEvent("client.stream.send.msg.start", trace.WithAttributes(
attribute.String("message.type", fmt.Sprintf("%T", m))),
)
err := cs.originalStream.SendMsg(m)
if err != nil {
select {
@ -61,10 +71,16 @@ func (cs *clientStream) SendMsg(m any) error {
case cs.finished <- err:
}
}
cs.span.AddEvent("client.stream.send.msg.finish", trace.WithAttributes(
attribute.String("message.type", fmt.Sprintf("%T", m))),
)
return err
}
func (cs *clientStream) RecvMsg(m any) error {
cs.span.AddEvent("client.stream.receive.msg.start", trace.WithAttributes(
attribute.String("message.type", fmt.Sprintf("%T", m))),
)
err := cs.originalStream.RecvMsg(m)
if err != nil || !cs.desc.ServerStreams {
select {
@ -72,5 +88,8 @@ func (cs *clientStream) RecvMsg(m any) error {
case cs.finished <- err:
}
}
cs.span.AddEvent("client.stream.receive.msg.finish", trace.WithAttributes(
attribute.String("message.type", fmt.Sprintf("%T", m))),
)
return err
}

View file

@ -50,7 +50,7 @@ func NewStreamClientInterceptor() grpc.StreamClientInterceptor {
finished := make(chan error)
done := make(chan struct{})
strWrp := newgRPCClientStream(str, desc, finished, done)
strWrp := newgRPCClientStream(str, desc, span, finished, done)
go func() {
defer close(done)
@ -107,7 +107,7 @@ func NewStreamServerInterceptor() grpc.StreamServerInterceptor {
trace.WithSpanKind(trace.SpanKindServer))
defer span.End()
err := handler(srv, newgRPCServerStream(ctx, ss))
err := handler(srv, newgRPCServerStream(ctx, ss, span))
setGRPCSpanStatus(span, err)
return err

View file

@ -2,7 +2,10 @@ package grpc
import (
"context"
"fmt"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/trace"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
)
@ -10,12 +13,14 @@ import (
type serverStream struct {
originalStream grpc.ServerStream
ctx context.Context // nolint:containedctx
span trace.Span
}
func newgRPCServerStream(ctx context.Context, originalStream grpc.ServerStream) grpc.ServerStream {
func newgRPCServerStream(ctx context.Context, originalStream grpc.ServerStream, span trace.Span) grpc.ServerStream {
return &serverStream{
originalStream: originalStream,
ctx: ctx,
span: span,
}
}
@ -36,9 +41,23 @@ func (ss *serverStream) Context() context.Context {
}
func (ss *serverStream) SendMsg(m any) error {
return ss.originalStream.SendMsg(m)
ss.span.AddEvent("server.stream.send.msg.start", trace.WithAttributes(
attribute.String("message.type", fmt.Sprintf("%T", m))),
)
err := ss.originalStream.SendMsg(m)
ss.span.AddEvent("server.stream.send.msg.finish", trace.WithAttributes(
attribute.String("message.type", fmt.Sprintf("%T", m))),
)
return err
}
func (ss *serverStream) RecvMsg(m any) error {
return ss.originalStream.RecvMsg(m)
ss.span.AddEvent("server.stream.receive.msg.start", trace.WithAttributes(
attribute.String("message.type", fmt.Sprintf("%T", m))),
)
err := ss.originalStream.RecvMsg(m)
ss.span.AddEvent("server.stream.receive.msg.finish", trace.WithAttributes(
attribute.String("message.type", fmt.Sprintf("%T", m))),
)
return err
}