forked from TrueCloudLab/frostfs-qos
96 lines
2.7 KiB
Go
96 lines
2.7 KiB
Go
package tagging
|
|
|
|
import (
|
|
"context"
|
|
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/metadata"
|
|
)
|
|
|
|
const (
|
|
ioTagHeader = "x-frostfs-io-tag"
|
|
)
|
|
|
|
// NewUnaryClientInteceptor creates new gRPC unary interceptor to set an IO tag to gRPC metadata.
|
|
func NewUnaryClientInteceptor() grpc.UnaryClientInterceptor {
|
|
return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
|
|
return invoker(setIOTagToGRPCMetadata(ctx), method, req, reply, cc, opts...)
|
|
}
|
|
}
|
|
|
|
// NewStreamClientInterceptor creates new gRPC stream interceptor to set an IO tag to gRPC metadata.
|
|
func NewStreamClientInterceptor() grpc.StreamClientInterceptor {
|
|
return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
|
|
return streamer(setIOTagToGRPCMetadata(ctx), desc, cc, method, opts...)
|
|
}
|
|
}
|
|
|
|
// NewUnaryServerInterceptor creates new gRPC unary interceptor to extract an IO tag to gRPC metadata.
|
|
func NewUnaryServerInterceptor() grpc.UnaryServerInterceptor {
|
|
return func(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
|
|
return handler(extractIOTagFromGRPCMetadata(ctx), req)
|
|
}
|
|
}
|
|
|
|
// NewStreamServerInterceptor creates new gRPC stream interceptor to extract an IO tag to gRPC metadata.
|
|
func NewStreamServerInterceptor() grpc.StreamServerInterceptor {
|
|
return func(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
|
|
return handler(srv, &serverStream{origin: ss})
|
|
}
|
|
}
|
|
|
|
func setIOTagToGRPCMetadata(ctx context.Context) context.Context {
|
|
ioTag, ok := IOTagFromContext(ctx)
|
|
if !ok {
|
|
return ctx
|
|
}
|
|
md, ok := metadata.FromOutgoingContext(ctx)
|
|
if !ok {
|
|
md = metadata.MD{}
|
|
}
|
|
md.Set(ioTagHeader, ioTag)
|
|
return metadata.NewOutgoingContext(ctx, md)
|
|
}
|
|
|
|
func extractIOTagFromGRPCMetadata(ctx context.Context) context.Context {
|
|
md, ok := metadata.FromIncomingContext(ctx)
|
|
if !ok {
|
|
return ctx
|
|
}
|
|
|
|
values := md.Get(ioTagHeader)
|
|
if len(values) > 0 {
|
|
return ContextWithIOTag(ctx, values[0])
|
|
}
|
|
return ctx
|
|
}
|
|
|
|
var _ grpc.ServerStream = &serverStream{}
|
|
|
|
type serverStream struct {
|
|
origin grpc.ServerStream
|
|
}
|
|
|
|
func (s *serverStream) Context() context.Context {
|
|
return extractIOTagFromGRPCMetadata(s.origin.Context())
|
|
}
|
|
|
|
func (s *serverStream) RecvMsg(m any) error {
|
|
return s.origin.RecvMsg(m)
|
|
}
|
|
|
|
func (s *serverStream) SendHeader(md metadata.MD) error {
|
|
return s.origin.SendHeader(md)
|
|
}
|
|
|
|
func (s *serverStream) SendMsg(m any) error {
|
|
return s.origin.SendMsg(m)
|
|
}
|
|
|
|
func (s *serverStream) SetHeader(md metadata.MD) error {
|
|
return s.origin.SetHeader(md)
|
|
}
|
|
|
|
func (s *serverStream) SetTrailer(md metadata.MD) {
|
|
s.origin.SetTrailer(md)
|
|
}
|