feat: add partition ring to dataobj consumer

Author: grobinson-grafana

docs/sources/shared/configuration.md

@@ -1252,6 +1252,66 @@ dataobj:
       # CLI flag: -dataobj-consumer.lifecycler.ID
       [id: <string> | default = "<hostname>"]
 
+    partition_ring:
+      # The key-value store used to share the hash ring across multiple
+      # instances. This option needs be set on ingesters, distributors,
+      # queriers, and rulers when running in microservices mode.
+      kvstore:
+        # Backend storage to use for the ring. Supported values are: consul,
+        # etcd, inmemory, memberlist, multi.
+        # CLI flag: -dataobj-consumer.partition-ring.store
+        [store: <string> | default = "memberlist"]
+
+        # The prefix for the keys in the store. Should end with a /.
+        # CLI flag: -dataobj-consumer.partition-ring.prefix
+        [prefix: <string> | default = "collectors/"]
+
+        # Configuration for a Consul client. Only applies if the selected
+        # kvstore is consul.
+        # The CLI flags prefix for this block configuration is:
+        # dataobj-consumer.partition-ring
+        [consul: <consul>]
+
+        # Configuration for an ETCD v3 client. Only applies if the selected
+        # kvstore is etcd.
+        # The CLI flags prefix for this block configuration is:
+        # dataobj-consumer.partition-ring
+        [etcd: <etcd>]
+
+        multi:
+          # Primary backend storage used by multi-client.
+          # CLI flag: -dataobj-consumer.partition-ring.multi.primary
+          [primary: <string> | default = ""]
+
+          # Secondary backend storage used by multi-client.
+          # CLI flag: -dataobj-consumer.partition-ring.multi.secondary
+          [secondary: <string> | default = ""]
+
+          # Mirror writes to secondary store.
+          # CLI flag: -dataobj-consumer.partition-ring.multi.mirror-enabled
+          [mirror_enabled: <boolean> | default = false]
+
+          # Timeout for storing value to secondary store.
+          # CLI flag: -dataobj-consumer.partition-ring.multi.mirror-timeout
+          [mirror_timeout: <duration> | default = 2s]
+
+      # Minimum number of owners to wait before a PENDING partition gets
+      # switched to ACTIVE.
+      # CLI flag: -dataobj-consumer.partition-ring.min-partition-owners-count
+      [min_partition_owners_count: <int> | default = 1]
+
+      # How long the minimum number of owners are enforced before a PENDING
+      # partition gets switched to ACTIVE.
+      # CLI flag: -dataobj-consumer.partition-ring.min-partition-owners-duration
+      [min_partition_owners_duration: <duration> | default = 10s]
+
+      # How long to wait before an INACTIVE partition is eligible for deletion.
+      # The partition is deleted only if it has been in INACTIVE state for at
+      # least the configured duration and it has no owners registered. A value
+      # of 0 disables partitions deletion.
+      # CLI flag: -dataobj-consumer.partition-ring.delete-inactive-partition-after
+      [delete_inactive_partition_after: <duration> | default = 13h]
+
     uploader:
       # The size of the SHA prefix to use for generating object storage keys for
       # data objects.
@@ -2824,6 +2884,7 @@ Configuration for a Consul client. Only applies if the selected kvstore is `cons
 - `common.storage.ring`
 - `compactor.ring`
 - `dataobj-consumer`
+- `dataobj-consumer.partition-ring`
 - `distributor.ring`
 - `index-gateway.ring`
 - `ingest-limits`
@@ -3075,6 +3136,7 @@ Configuration for an ETCD v3 client. Only applies if the selected kvstore is `et
 - `common.storage.ring`
 - `compactor.ring`
 - `dataobj-consumer`
+- `dataobj-consumer.partition-ring`
 - `distributor.ring`
 - `index-gateway.ring`
 - `ingest-limits`
@@ -7498,6 +7560,7 @@ The TLS configuration. The supported CLI flags `<prefix>` used to reference this
 - `compactor.grpc-client`
 - `compactor.ring.etcd`
 - `dataobj-consumer.etcd`
+- `dataobj-consumer.partition-ring.etcd`
 - `distributor.ring.etcd`
 - `etcd`
 - `frontend.grpc-client-config`

pkg/dataobj/consumer/config.go

@@ -8,14 +8,16 @@ import (
 
 	"github.com/grafana/loki/v3/pkg/dataobj/consumer/logsobj"
 	"github.com/grafana/loki/v3/pkg/dataobj/uploader"
+	"github.com/grafana/loki/v3/pkg/kafka/partitionring"
 	util_log "github.com/grafana/loki/v3/pkg/util/log"
 )
 
 type Config struct {
-	BuilderConfig    logsobj.BuilderConfig `yaml:"builder,omitempty"`
-	LifecyclerConfig ring.LifecyclerConfig `yaml:"lifecycler,omitempty"`
-	UploaderConfig   uploader.Config       `yaml:"uploader"`
-	IdleFlushTimeout time.Duration         `yaml:"idle_flush_timeout"`
+	BuilderConfig       logsobj.BuilderConfig `yaml:"builder,omitempty"`
+	LifecyclerConfig    ring.LifecyclerConfig `yaml:"lifecycler,omitempty"`
+	PartitionRingConfig partitionring.Config  `yaml:"partition_ring" category:"experimental"`
+	UploaderConfig      uploader.Config       `yaml:"uploader"`
+	IdleFlushTimeout    time.Duration         `yaml:"idle_flush_timeout"`
 }
 
 func (cfg *Config) Validate() error {
@@ -38,6 +40,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 func (cfg *Config) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
 	cfg.BuilderConfig.RegisterFlagsWithPrefix(prefix, f)
 	cfg.LifecyclerConfig.RegisterFlagsWithPrefix(prefix, f, util_log.Logger)
+	cfg.PartitionRingConfig.RegisterFlagsWithPrefix(prefix, f)
 	cfg.UploaderConfig.RegisterFlagsWithPrefix(prefix, f)
 
 	f.DurationVar(&cfg.IdleFlushTimeout, prefix+"idle-flush-timeout", 60*60*time.Second, "The maximum amount of time to wait in seconds before flushing an object that is no longer receiving new writes")

pkg/dataobj/consumer/service.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
+	"github.com/grafana/dskit/kv"
 	"github.com/grafana/dskit/ring"
 	"github.com/grafana/dskit/services"
 	"github.com/prometheus/client_golang/prometheus"
@@ -15,22 +16,26 @@ import (
 	"github.com/grafana/loki/v3/pkg/dataobj/metastore"
 	"github.com/grafana/loki/v3/pkg/kafka"
 	"github.com/grafana/loki/v3/pkg/kafka/client"
+	"github.com/grafana/loki/v3/pkg/kafka/partitionring"
 	"github.com/grafana/loki/v3/pkg/scratch"
 )
 
 const (
-	RingKey  = "dataobj-consumer"
-	RingName = "dataobj-consumer"
+	RingKey           = "dataobj-consumer"
+	RingName          = "dataobj-consumer"
+	PartitionRingKey  = "dataobj-consumer-partitions-key"
+	PartitionRingName = "dataobj-consumer-partitions"
 )
 
 type Service struct {
 	services.Service
-	cfg             Config
-	metastoreEvents *kgo.Client
-	lifecycler      *ring.Lifecycler
-	watcher         *services.FailureWatcher
-	logger          log.Logger
-	reg             prometheus.Registerer
+	cfg                         Config
+	metastoreEvents             *kgo.Client
+	lifecycler                  *ring.Lifecycler
+	partitionInstanceLifecycler *ring.PartitionInstanceLifecycler
+	watcher                     *services.FailureWatcher
+	logger                      log.Logger
+	reg                         prometheus.Registerer
 }
 
 func New(kafkaCfg kafka.Config, cfg Config, _ metastore.Config, _ objstore.Bucket, _ scratch.Store, _ string, _ ring.PartitionRingReader, reg prometheus.Registerer, logger log.Logger) (*Service, error) {
@@ -69,20 +74,57 @@ func New(kafkaCfg kafka.Config, cfg Config, _ metastore.Config, _ objstore.Bucke
 	}
 	s.lifecycler = lifecycler
 
+	// An instance must register itself in the partition ring. Each instance
+	// is responsible for consuming exactly one partition determined by
+	// its partition ID. Once ready, the instance will declare its partition
+	// as active in the partition ring. This is how distributors know which
+	// partitions have a ready consumer.
+	partitionID, err := partitionring.ExtractPartitionID(cfg.LifecyclerConfig.ID)
+	if err != nil {
+		return nil, fmt.Errorf("failed to extract partition ID from lifecycler configuration: %w", err)
+	}
+	partitionRingKV := cfg.PartitionRingConfig.KVStore.Mock
+	// The mock KV is used in tests. If this is not a test then we must
+	// initialize a real kv.
+	if partitionRingKV == nil {
+		partitionRingKV, err = kv.NewClient(
+			cfg.PartitionRingConfig.KVStore,
+			ring.GetPartitionRingCodec(),
+			kv.RegistererWithKVName(reg, "dataobj-consumer-lifecycler"),
+			logger,
+		)
+		if err != nil {
+			return nil, fmt.Errorf("failed to set up partition ring: %w", err)
+		}
+	}
+	partitionInstanceLifecycler := ring.NewPartitionInstanceLifecycler(
+		cfg.PartitionRingConfig.ToLifecyclerConfig(partitionID, cfg.LifecyclerConfig.ID),
+		PartitionRingName,
+		PartitionRingKey,
+		partitionRingKV,
+		logger,
+		prometheus.WrapRegistererWithPrefix("loki_", reg))
+	s.partitionInstanceLifecycler = partitionInstanceLifecycler
+
 	watcher := services.NewFailureWatcher()
 	watcher.WatchService(lifecycler)
+	watcher.WatchService(partitionInstanceLifecycler)
 	s.watcher = watcher
 
 	s.Service = services.NewBasicService(s.starting, s.running, s.stopping)
 	return s, nil
 }
 
 // starting implements the Service interface's starting method.
-func (s *Service) starting(ctx context.Context) (err error) {
+func (s *Service) starting(ctx context.Context) error {
 	level.Info(s.logger).Log("msg", "starting")
 	if err := services.StartAndAwaitRunning(ctx, s.lifecycler); err != nil {
 		return fmt.Errorf("failed to start lifecycler: %w", err)
 	}
+	if err := services.StartAndAwaitRunning(ctx, s.partitionInstanceLifecycler); err != nil {
+		return fmt.Errorf("failed to start partition instance lifecycler: %w", err)
+	}
+
 	return nil
 }
 
@@ -96,6 +138,9 @@ func (s *Service) running(ctx context.Context) error {
 func (s *Service) stopping(failureCase error) error {
 	level.Info(s.logger).Log("msg", "stopping")
 	ctx := context.TODO()
+	if err := services.StopAndAwaitTerminated(ctx, s.partitionInstanceLifecycler); err != nil {
+		level.Warn(s.logger).Log("msg", "failed to stop partition instance lifecycler", "err", err)
+	}
 	if err := services.StopAndAwaitTerminated(ctx, s.lifecycler); err != nil {
 		level.Warn(s.logger).Log("msg", "failed to stop lifecycler", "err", err)
 	}

pkg/loki/config_wrapper.go

@@ -379,6 +379,7 @@ func applyConfigToRings(r, defaults *ConfigWrapper, rc lokiring.RingConfig, merg
 		r.DataObj.Consumer.LifecyclerConfig.ListenPort = rc.ListenPort
 		r.DataObj.Consumer.LifecyclerConfig.ObservePeriod = rc.ObservePeriod
 		r.DataObj.Consumer.LifecyclerConfig.EnableInet6 = rc.EnableIPv6
+		r.DataObj.Consumer.PartitionRingConfig.KVStore = rc.KVStore
 	}
 }