frostfs-observability/tracing/grpc/server.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
}