frostfs-node/internal/qos/grpc.go

package qos

import (
	"context"

	"git.frostfs.info/TrueCloudLab/frostfs-qos/limiting"
	"git.frostfs.info/TrueCloudLab/frostfs-qos/tagging"
	apistatus "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/client/status"
	"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.IsLocal() {
			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.IsLocal() {
			tag = IOTagInternal
		}
		ctx = tagging.ContextWithIOTag(ctx, tag.String())
		return streamer(ctx, desc, cc, method, opts...)
	}
}

func NewMaxActiveRPCLimiterUnaryServerInterceptor(getLimiter func() limiting.Limiter) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
		if tag, ok := tagging.IOTagFromContext(ctx); ok && tag == IOTagCritical.String() {
			return handler(ctx, req)
		}

		release, ok := getLimiter().Acquire(info.FullMethod)
		if !ok {
			return nil, new(apistatus.ResourceExhausted)
		}
		defer release()

		return handler(ctx, req)
	}
}

//nolint:contextcheck (grpc.ServerStream manages the context itself)
func NewMaxActiveRPCLimiterStreamServerInterceptor(getLimiter func() limiting.Limiter) grpc.StreamServerInterceptor {
	return func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		if tag, ok := tagging.IOTagFromContext(ss.Context()); ok && tag == IOTagCritical.String() {
			return handler(srv, ss)
		}

		release, ok := getLimiter().Acquire(info.FullMethod)
		if !ok {
			return new(apistatus.ResourceExhausted)
		}
		defer release()

		return handler(srv, ss)
	}
}