frostfs-node/internal/qos/grpc.go

package qos

import (
	"context"

	"git.frostfs.info/TrueCloudLab/frostfs-qos/tagging"
	"google.golang.org/grpc"
)

func NewSetCriticalIOTagUnaryServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
		ctx = tagging.ContextWithIOTag(ctx, IOTagCritical.String())
		return handler(ctx, req)
	}
}

func NewAdjustOutgoingIOTagUnaryClientInterceptor() grpc.UnaryClientInterceptor {
	return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		rawTag, ok := tagging.IOTagFromContext(ctx)
		if !ok {
			return invoker(ctx, method, req, reply, cc, opts...)
		}
		tag, err := FromRawString(rawTag)
		if err != nil {
			tag = IOTagClient
		}
		if tag == IOTagBackground || tag == IOTagPolicer || tag == IOTagWritecache {
			tag = IOTagInternal
		}
		ctx = tagging.ContextWithIOTag(ctx, tag.String())
		return invoker(ctx, method, req, reply, cc, opts...)
	}
}

func NewAdjustOutgoingIOTagStreamClientInterceptor() grpc.StreamClientInterceptor {
	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
		rawTag, ok := tagging.IOTagFromContext(ctx)
		if !ok {
			return streamer(ctx, desc, cc, method, opts...)
		}
		tag, err := FromRawString(rawTag)
		if err != nil {
			tag = IOTagClient
		}
		if tag == IOTagBackground || tag == IOTagPolicer || tag == IOTagWritecache {
			tag = IOTagInternal
		}
		ctx = tagging.ContextWithIOTag(ctx, tag.String())
		return streamer(ctx, desc, cc, method, opts...)
	}
}