63 lines
1.6 KiB
Go
63 lines
1.6 KiB
Go
package grpc
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"go.opentelemetry.io/otel/attribute"
|
|
"go.opentelemetry.io/otel/trace"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/metadata"
|
|
)
|
|
|
|
type serverStream struct {
|
|
originalStream grpc.ServerStream
|
|
ctx context.Context // nolint:containedctx
|
|
span trace.Span
|
|
}
|
|
|
|
func newgRPCServerStream(ctx context.Context, originalStream grpc.ServerStream, span trace.Span) grpc.ServerStream {
|
|
return &serverStream{
|
|
originalStream: originalStream,
|
|
ctx: ctx,
|
|
span: span,
|
|
}
|
|
}
|
|
|
|
func (ss *serverStream) SetHeader(md metadata.MD) error {
|
|
return ss.originalStream.SendHeader(md)
|
|
}
|
|
|
|
func (ss *serverStream) SendHeader(md metadata.MD) error {
|
|
return ss.originalStream.SendHeader(md)
|
|
}
|
|
|
|
func (ss *serverStream) SetTrailer(md metadata.MD) {
|
|
ss.originalStream.SetTrailer(md)
|
|
}
|
|
|
|
func (ss *serverStream) Context() context.Context {
|
|
return ss.ctx
|
|
}
|
|
|
|
func (ss *serverStream) SendMsg(m any) error {
|
|
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 {
|
|
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
|
|
}
|