feat: add limits/frontend pkg (#16658)

Author: grobinson-grafana

pkg/limits/frontend/client/client.go

@@ -0,0 +1,163 @@
+// Package client provides gRPC client implementation for limits-frontend.
+// An example use case is the distributor, which needs to ask limits-frontend
+// if a push request has exceeded per-tenant limits.
+package client
+
+import (
+	"flag"
+	"io"
+	"time"
+
+	"github.com/go-kit/log"
+	"github.com/grafana/dskit/grpcclient"
+	"github.com/grafana/dskit/middleware"
+	"github.com/grafana/dskit/ring"
+	ring_client "github.com/grafana/dskit/ring/client"
+	"github.com/grpc-ecosystem/grpc-opentracing/go/otgrpc"
+	"github.com/opentracing/opentracing-go"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/health/grpc_health_v1"
+
+	"github.com/grafana/loki/v3/pkg/logproto"
+	"github.com/grafana/loki/v3/pkg/util/server"
+)
+
+var (
+	LimitsRead = ring.NewOp([]ring.InstanceState{ring.ACTIVE}, nil)
+)
+
+var (
+	frontendClients = prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: "loki_ingest_limits_frontend_clients",
+		Help: "The current number of ingest limits frontend clients.",
+	})
+	frontendRequestDuration = promauto.NewHistogramVec(prometheus.HistogramOpts{
+		Name:    "loki_ingest_limits_frontend_client_request_duration_seconds",
+		Help:    "Time spent doing ingest limits frontend requests.",
+		Buckets: prometheus.ExponentialBuckets(0.001, 4, 6),
+	}, []string{"operation", "status_code"})
+)
+
+// Config contains the config for an ingest-limits-frontend client.
+type Config struct {
+	GRPCClientConfig             grpcclient.Config              `yaml:"grpc_client_config" doc:"description=Configures client gRPC connections to limits service."`
+	PoolConfig                   PoolConfig                     `yaml:"pool_config,omitempty" doc:"description=Configures client gRPC connections pool to limits service."`
+	GRPCUnaryClientInterceptors  []grpc.UnaryClientInterceptor  `yaml:"-"`
+	GRCPStreamClientInterceptors []grpc.StreamClientInterceptor `yaml:"-"`
+
+	// Internal is used to indicate that this client communicates on behalf of
+	// a machine and not a user. When Internal = true, the client won't attempt
+	// to inject an userid into the context.
+	Internal bool `yaml:"-"`
+}
+
+func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
+	cfg.RegisterFlagsWithPrefix("ingest-limits-frontend-client", f)
+}
+
+func (cfg *Config) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
+	cfg.GRPCClientConfig.RegisterFlagsWithPrefix(prefix, f)
+	cfg.PoolConfig.RegisterFlagsWithPrefix(prefix, f)
+}
+
+// PoolConfig contains the config for a pool of ingest-limits-frontend clients.
+type PoolConfig struct {
+	ClientCleanupPeriod     time.Duration `yaml:"client_cleanup_period"`
+	HealthCheckIngestLimits bool          `yaml:"health_check_ingest_limits"`
+	RemoteTimeout           time.Duration `yaml:"remote_timeout"`
+}
+
+func (cfg *PoolConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
+	f.DurationVar(&cfg.ClientCleanupPeriod, prefix+".client-cleanup-period", 15*time.Second, "How frequently to clean up clients for ingest-limits-frontend that have gone away.")
+	f.BoolVar(&cfg.HealthCheckIngestLimits, prefix+".health-check-ingest-limits", true, "Run a health check on each ingest-limits-frontend client during periodic cleanup.")
+	f.DurationVar(&cfg.RemoteTimeout, prefix+".remote-timeout", 1*time.Second, "Timeout for the health check.")
+}
+
+// Client is a gRPC client for the ingest-limits-frontend.
+type Client struct {
+	logproto.IngestLimitsFrontendClient
+	grpc_health_v1.HealthClient
+	io.Closer
+}
+
+// NewClient returns a new Client for the specified ingest-limits-frontend.
+func NewClient(cfg Config, addr string) (*Client, error) {
+	opts := []grpc.DialOption{
+		grpc.WithDefaultCallOptions(cfg.GRPCClientConfig.CallOptions()...),
+	}
+	dialOpts, err := cfg.GRPCClientConfig.DialOption(getGRPCInterceptors(&cfg))
+	if err != nil {
+		return nil, err
+	}
+	opts = append(opts, dialOpts...)
+	// nolint:staticcheck // grpc.Dial() has been deprecated; we'll address it before upgrading to gRPC 2.
+	conn, err := grpc.Dial(addr, opts...)
+	if err != nil {
+		return nil, err
+	}
+	return &Client{
+		IngestLimitsFrontendClient: logproto.NewIngestLimitsFrontendClient(conn),
+		HealthClient:               grpc_health_v1.NewHealthClient(conn),
+		Closer:                     conn,
+	}, nil
+}
+
+// getInterceptors returns the gRPC interceptors for the given Config.
+func getGRPCInterceptors(cfg *Config) ([]grpc.UnaryClientInterceptor, []grpc.StreamClientInterceptor) {
+	var (
+		unaryInterceptors  []grpc.UnaryClientInterceptor
+		streamInterceptors []grpc.StreamClientInterceptor
+	)
+
+	unaryInterceptors = append(unaryInterceptors, cfg.GRPCUnaryClientInterceptors...)
+	unaryInterceptors = append(unaryInterceptors, server.UnaryClientQueryTagsInterceptor)
+	unaryInterceptors = append(unaryInterceptors, server.UnaryClientHTTPHeadersInterceptor)
+	unaryInterceptors = append(unaryInterceptors, otgrpc.OpenTracingClientInterceptor(opentracing.GlobalTracer()))
+	if !cfg.Internal {
+		unaryInterceptors = append(unaryInterceptors, middleware.ClientUserHeaderInterceptor)
+	}
+	unaryInterceptors = append(unaryInterceptors, middleware.UnaryClientInstrumentInterceptor(frontendRequestDuration))
+
+	streamInterceptors = append(streamInterceptors, cfg.GRCPStreamClientInterceptors...)
+	streamInterceptors = append(streamInterceptors, server.StreamClientQueryTagsInterceptor)
+	streamInterceptors = append(streamInterceptors, server.StreamClientHTTPHeadersInterceptor)
+	streamInterceptors = append(streamInterceptors, otgrpc.OpenTracingStreamClientInterceptor(opentracing.GlobalTracer()))
+	if !cfg.Internal {
+		streamInterceptors = append(streamInterceptors, middleware.StreamClientUserHeaderInterceptor)
+	}
+	streamInterceptors = append(streamInterceptors, middleware.StreamClientInstrumentInterceptor(frontendRequestDuration))
+
+	return unaryInterceptors, streamInterceptors
+}
+
+// NewPool returns a new pool of clients for the ingest-limits-frontend.
+func NewPool(
+	name string,
+	cfg PoolConfig,
+	ring ring.ReadRing,
+	factory ring_client.PoolFactory,
+	logger log.Logger,
+) *ring_client.Pool {
+	poolCfg := ring_client.PoolConfig{
+		CheckInterval:      cfg.ClientCleanupPeriod,
+		HealthCheckEnabled: cfg.HealthCheckIngestLimits,
+		HealthCheckTimeout: cfg.RemoteTimeout,
+	}
+	return ring_client.NewPool(
+		name,
+		poolCfg,
+		ring_client.NewRingServiceDiscovery(ring),
+		factory,
+		frontendClients,
+		logger,
+	)
+}
+
+// NewPoolFactory returns a new factory for ingest-limits-frontend clients.
+func NewPoolFactory(cfg Config) ring_client.PoolFactory {
+	return ring_client.PoolAddrFunc(func(addr string) (ring_client.PoolClient, error) {
+		return NewClient(cfg, addr)
+	})
+}

pkg/limits/frontend/frontend.go

@@ -0,0 +1,208 @@
+// Package frontend contains provides a frontend service for ingest limits.
+// It is responsible for receiving and answering gRPC requests from distributors,
+// such as exceeds limits requests, forwarding them to individual limits backends,
+// gathering and aggregating their responses (where required), and returning
+// the final result.
+package frontend
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"net/http"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+	"github.com/grafana/dskit/ring"
+	"github.com/grafana/dskit/services"
+	"github.com/grafana/dskit/user"
+	"github.com/prometheus/client_golang/prometheus"
+
+	limits_client "github.com/grafana/loki/v3/pkg/limits/client"
+	"github.com/grafana/loki/v3/pkg/logproto"
+	"github.com/grafana/loki/v3/pkg/util"
+	util_log "github.com/grafana/loki/v3/pkg/util/log"
+)
+
+const (
+	RingKey  = "ingest-limits-frontend"
+	RingName = "ingest-limits-frontend"
+)
+
+// Config contains the config for an ingest-limits-frontend.
+type Config struct {
+	ClientConfig     limits_client.Config  `yaml:"client_config"`
+	LifecyclerConfig ring.LifecyclerConfig `yaml:"lifecycler,omitempty"`
+}
+
+func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
+	cfg.ClientConfig.RegisterFlagsWithPrefix("ingest-limits-frontend", f)
+	cfg.LifecyclerConfig.RegisterFlagsWithPrefix("ingest-limits-frontend.", f, util_log.Logger)
+}
+
+func (cfg *Config) Validate() error {
+	if err := cfg.ClientConfig.GRPCClientConfig.Validate(); err != nil {
+		return fmt.Errorf("invalid gRPC client config: %w", err)
+	}
+	return nil
+}
+
+// Frontend is the limits-frontend service, and acts a service wrapper for
+// all components needed to run the limits-frontend.
+type Frontend struct {
+	services.Service
+
+	cfg    Config
+	logger log.Logger
+
+	subservices        *services.Manager
+	subservicesWatcher *services.FailureWatcher
+
+	limits IngestLimitsService
+
+	lifecycler        *ring.Lifecycler
+	lifecyclerWatcher *services.FailureWatcher
+}
+
+// New returns a new Frontend.
+func New(cfg Config, ringName string, limitsRing ring.ReadRing, limits Limits, logger log.Logger, reg prometheus.Registerer) (*Frontend, error) {
+	var servs []services.Service
+
+	factory := limits_client.NewPoolFactory(cfg.ClientConfig)
+	pool := limits_client.NewPool(ringName, cfg.ClientConfig.PoolConfig, limitsRing, factory, logger)
+	limitsSrv := NewRingIngestLimitsService(limitsRing, pool, limits, logger, reg)
+
+	f := &Frontend{
+		cfg:    cfg,
+		logger: logger,
+		limits: limitsSrv,
+	}
+
+	var err error
+	f.lifecycler, err = ring.NewLifecycler(cfg.LifecyclerConfig, f, RingName, RingKey, true, logger, reg)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create %s lifecycler: %w", RingName, err)
+	}
+	// Watch the lifecycler
+	f.lifecyclerWatcher = services.NewFailureWatcher()
+	f.lifecyclerWatcher.WatchService(f.lifecycler)
+
+	servs = append(servs, f.lifecycler)
+	servs = append(servs, pool)
+	mgr, err := services.NewManager(servs...)
+	if err != nil {
+		return nil, err
+	}
+
+	f.subservices = mgr
+	f.subservicesWatcher = services.NewFailureWatcher()
+	f.subservicesWatcher.WatchManager(f.subservices)
+	f.Service = services.NewBasicService(f.starting, f.running, f.stopping)
+
+	return f, nil
+}
+
+// Flush implements ring.FlushTransferer. It transfers state to another ingest limits frontend instance.
+func (f *Frontend) Flush() {}
+
+// TransferOut implements ring.FlushTransferer. It transfers state to another ingest limits frontend instance.
+func (f *Frontend) TransferOut(_ context.Context) error {
+	return nil
+}
+
+// starting implements services.Service.
+func (f *Frontend) starting(ctx context.Context) (err error) {
+	defer func() {
+		if err == nil {
+			return
+		}
+		stopErr := services.StopManagerAndAwaitStopped(context.Background(), f.subservices)
+		if stopErr != nil {
+			level.Error(f.logger).Log("msg", "failed to stop subservices", "err", stopErr)
+		}
+	}()
+
+	level.Info(f.logger).Log("msg", "starting subservices")
+	if err := services.StartManagerAndAwaitHealthy(ctx, f.subservices); err != nil {
+		return fmt.Errorf("failed to start subservices: %w", err)
+	}
+
+	return nil
+}
+
+// running implements services.Service.
+func (f *Frontend) running(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return nil
+	case err := <-f.subservicesWatcher.Chan():
+		return fmt.Errorf("ingest limits frontend subservice failed: %w", err)
+	}
+}
+
+// stopping implements services.Service.
+func (f *Frontend) stopping(_ error) error {
+	return services.StopManagerAndAwaitStopped(context.Background(), f.subservices)
+}
+
+// ExceedsLimits implements logproto.IngestLimitsFrontendClient.
+func (f *Frontend) ExceedsLimits(ctx context.Context, r *logproto.ExceedsLimitsRequest) (*logproto.ExceedsLimitsResponse, error) {
+	return f.limits.ExceedsLimits(ctx, r)
+}
+
+type exceedsLimitsRequest struct {
+	TenantID     string   `json:"tenantID"`
+	StreamHashes []uint64 `json:"streamHashes"`
+}
+
+type exceedsLimitsResponse struct {
+	RejectedStreams []*logproto.RejectedStream `json:"rejectedStreams,omitempty"`
+}
+
+// ServeHTTP implements http.Handler.
+func (f *Frontend) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	var req exceedsLimitsRequest
+	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+		level.Error(f.logger).Log("msg", "error unmarshalling request body", "err", err)
+		http.Error(w, "error unmarshalling request body", http.StatusBadRequest)
+		return
+	}
+
+	if req.TenantID == "" {
+		http.Error(w, "tenantID is required", http.StatusBadRequest)
+		return
+	}
+
+	// Convert request to protobuf format
+	protoReq := &logproto.ExceedsLimitsRequest{
+		Tenant:  req.TenantID,
+		Streams: make([]*logproto.StreamMetadata, len(req.StreamHashes)),
+	}
+	for i, hash := range req.StreamHashes {
+		protoReq.Streams[i] = &logproto.StreamMetadata{
+			StreamHash: hash,
+		}
+	}
+
+	ctx, err := user.InjectIntoGRPCRequest(user.InjectOrgID(r.Context(), req.TenantID))
+	if err != nil {
+		http.Error(w, "failed to inject org ID", http.StatusInternalServerError)
+		return
+	}
+
+	// Call the service
+	resp, err := f.limits.ExceedsLimits(ctx, protoReq)
+	if err != nil {
+		level.Error(f.logger).Log("msg", "error checking limits", "err", err)
+		http.Error(w, "error checking limits", http.StatusInternalServerError)
+		return
+	}
+
+	// Convert response to JSON format
+	jsonResp := exceedsLimitsResponse{
+		RejectedStreams: resp.RejectedStreams,
+	}
+
+	util.WriteJSONResponse(w, jsonResp)
+}