feat: add limits package (#16489)

Author: grobinson-grafana

pkg/kafka/encoding.go

@@ -16,11 +16,17 @@ import (
 	"github.com/grafana/loki/v3/pkg/logql/syntax"
 )
 
-var encoderPool = sync.Pool{
-	New: func() any {
-		return &logproto.Stream{}
-	},
-}
+const (
+	metadataTopicSuffix = ".metadata"
+)
+
+var (
+	encoderPool = sync.Pool{
+		New: func() any {
+			return &logproto.Stream{}
+		},
+	}
+)
 
 // Encode converts a logproto.Stream into one or more Kafka records.
 // It handles splitting large streams into multiple records if necessary.
@@ -188,3 +194,45 @@ func (d *Decoder) DecodeWithoutLabels(data []byte) (logproto.Stream, error) {
 func sovPush(x uint64) (n int) {
 	return (math_bits.Len64(x|1) + 6) / 7
 }
+
+// EncodeStreamMetadata encodes the stream metadata into a Kafka record
+// using the tenantID as the key and partition as the target partition
+func EncodeStreamMetadata(partition int32, topic string, tenantID string, streamHash uint64) (*kgo.Record, error) {
+	// Validate stream hash
+	if streamHash == 0 {
+		return nil, fmt.Errorf("invalid stream hash '%d'", streamHash)
+	}
+
+	metadata := logproto.StreamMetadata{
+		StreamHash: streamHash,
+	}
+
+	// Encode the metadata into a byte slice
+	value, err := metadata.Marshal()
+	if err != nil {
+		return nil, err
+	}
+
+	return &kgo.Record{
+		Key:       []byte(tenantID),
+		Value:     value,
+		Partition: partition,
+		Topic:     MetadataTopicFor(topic),
+	}, nil
+}
+
+// DecodeStreamMetadata decodes a Kafka record into a StreamMetadata.
+// It returns the decoded metadata and any error encountered.
+func DecodeStreamMetadata(record *kgo.Record) (*logproto.StreamMetadata, error) {
+	var metadata logproto.StreamMetadata
+	if err := metadata.Unmarshal(record.Value); err != nil {
+		return nil, fmt.Errorf("failed to unmarshal stream metadata: %w", err)
+	}
+
+	return &metadata, nil
+}
+
+// MetadataTopicFor returns the metadata topic name for the given topic.
+func MetadataTopicFor(topic string) string {
+	return topic + metadataTopicSuffix
+}

pkg/kafka/encoding_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/stretchr/testify/require"
+	"github.com/twmb/franz-go/pkg/kgo"
 
 	"github.com/grafana/loki/v3/pkg/logproto"
 )
@@ -149,3 +150,66 @@ func generateRandomString(length int) string {
 	}
 	return string(b)
 }
+
+func TestEncodeDecodeStreamMetadata(t *testing.T) {
+	tests := []struct {
+		name      string
+		hash      uint64
+		partition int32
+		topic     string
+		tenantID  string
+		expectErr bool
+	}{
+		{
+			name:      "Valid metadata",
+			hash:      12345,
+			partition: 1,
+			topic:     "logs",
+			tenantID:  "tenant-1",
+			expectErr: false,
+		},
+		{
+			name:      "Zero hash - should error",
+			hash:      0,
+			partition: 3,
+			topic:     "traces",
+			tenantID:  "tenant-3",
+			expectErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Encode metadata
+			record, err := EncodeStreamMetadata(tt.partition, tt.topic, tt.tenantID, tt.hash)
+			if tt.expectErr {
+				require.Error(t, err)
+				require.Nil(t, record)
+				return
+			}
+
+			require.NotNil(t, record)
+			require.NotNil(t, record.Value)
+			require.Equal(t, tt.topic+metadataTopicSuffix, record.Topic)
+			require.Equal(t, tt.partition, record.Partition)
+			require.Equal(t, []byte(tt.tenantID), record.Key)
+
+			// Decode metadata
+			metadata, err := DecodeStreamMetadata(record)
+			require.NoError(t, err)
+			require.NotNil(t, metadata)
+
+			// Verify decoded values
+			require.Equal(t, tt.hash, metadata.StreamHash)
+		})
+	}
+
+	t.Run("Decode invalid value", func(t *testing.T) {
+		record := &kgo.Record{
+			Value: []byte("invalid data"),
+		}
+		metadata, err := DecodeStreamMetadata(record)
+		require.Error(t, err)
+		require.Nil(t, metadata)
+	})
+}

pkg/limits/client/client.go

@@ -0,0 +1,153 @@
+// Package client provides gRPC client implementation for limits service.
+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 (
+	backendClients = prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: "loki_ingest_limits_backend_clients",
+		Help: "The current number of ingest limits backend clients.",
+	})
+	backendRequestDuration = promauto.NewHistogramVec(prometheus.HistogramOpts{
+		Name:    "loki_ingest_limits_backend_client_request_duration_seconds",
+		Help:    "Time spent doing ingest limits backend requests.",
+		Buckets: prometheus.ExponentialBuckets(0.001, 4, 6),
+	}, []string{"operation", "status_code"})
+)
+
+// Config contains the config for an ingest-limits 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) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
+	cfg.GRPCClientConfig.RegisterFlagsWithPrefix(prefix+".limits-client", f)
+	cfg.PoolConfig.RegisterFlagsWithPrefix(prefix, f)
+}
+
+// PoolConfig contains the config for a pool of ingest-limits 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 that have gone away.")
+	f.BoolVar(&cfg.HealthCheckIngestLimits, prefix+".health-check-ingest-limits", true, "Run a health check on each ingest-limits 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.
+type Client struct {
+	logproto.IngestLimitsClient
+	grpc_health_v1.HealthClient
+	io.Closer
+}
+
+// NewClient returns a new Client for the specified ingest-limits.
+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{
+		IngestLimitsClient: logproto.NewIngestLimitsClient(conn),
+		HealthClient:       grpc_health_v1.NewHealthClient(conn),
+		Closer:             conn,
+	}, nil
+}
+
+// getInterceptors returns the gRPC interceptors for the given ClientConfig.
+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(backendRequestDuration))
+
+	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(backendRequestDuration))
+
+	return unaryInterceptors, streamInterceptors
+}
+
+// NewPool returns a new pool of clients for the ingest-limits.
+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,
+		backendClients,
+		logger,
+	)
+}
+
+// NewPoolFactory returns a new factory for ingest-limits clients.
+func NewPoolFactory(cfg Config) ring_client.PoolFactory {
+	return ring_client.PoolAddrFunc(func(addr string) (ring_client.PoolClient, error) {
+		return NewClient(cfg, addr)
+	})
+}