Loki: Modifies TableManager to use IndexGateway ring

docs/sources/configuration/_index.md

@@ -132,6 +132,9 @@ Pass the `-config.expand-env` flag at the command line to enable this way of set
 # key value store.
 [ingester: <ingester>]
 
+# Configures the index gateway server.
+[index_gateway: <index_gateway>]
+
 # Configures where Loki will store data.
 [storage_config: <storage_config>]
 
@@ -1757,6 +1760,10 @@ boltdb_shipper:
     # The CLI flags prefix for this block config is: boltdb.shipper.index-gateway-client
     [grpc_client_config: <grpc_client_config>]
 
+    # Configures if log gateway requests should log debug messages or not.
+    # CLI flag: -boltdb.shipper.index-gateway-client.log-gateway-requests
+    [log_gateway_requests: <bool> | default = false]
+
 # Cache validity for active index entries. Should be no higher than
 # the chunk_idle_period in the ingester settings.
 # CLI flag: -store.index-cache-validity
@@ -2397,6 +2404,26 @@ backoff_config:
   [max_retries: <int> | default = 10]
 ```
 
+## index_gateway
+
+The `index_gateway` block configures the Loki index gateway server, responsible for serving index queries
+without the need to constantly interact with the object store.
+
+```yaml
+# Defines in which mode the index gateway server will operate (default to 'simple').
+# It supports two modes:
+# 'simple': an index gateway server instance is responsible for handling,
+#     storing and returning requests for all indices for all tenants.
+# 'ring': an index gateway server instance is responsible for a subset of tenants instead
+#     of all tenants.
+[mode: <string> | default = simple]
+
+# Defines the ring to be used by the index gateway servers and clients in case the servers
+# are configured to run in 'ring' mode. In case this isn't configured, this block supports
+# inheriting configuration from the common ring section.
+[ring: <ring>]
+```
+
 ## table_manager
 
 The `table_manager` block configures the Loki table-manager.

pkg/loki/loki.go

@@ -253,7 +253,7 @@ type Loki struct {
 	QueryFrontEndTripperware basetripper.Tripperware
 	queryScheduler           *scheduler.Scheduler
 	usageReport              *usagestats.Reporter
-	indexGatewayRing         *ring.Ring
+	indexGatewayRingManager  *indexgateway.RingManager
 
 	clientMetrics storage.ClientMetrics
 
@@ -512,7 +512,7 @@ func (t *Loki) setupModuleManager() error {
 		Ruler:                    {Ring, Server, Store, RulerStorage, IngesterQuerier, Overrides, TenantConfigs, UsageReport},
 		TableManager:             {Server, UsageReport},
 		Compactor:                {Server, Overrides, MemberlistKV, UsageReport},
-		IndexGateway:             {Server, Store, Overrides, UsageReport, MemberlistKV},
+		IndexGateway:             {Server, Store, Overrides, UsageReport, MemberlistKV, IndexGatewayRing},
 		IngesterQuerier:          {Ring},
 		IndexGatewayRing:         {RuntimeConfig, Server, MemberlistKV},
 		All:                      {QueryScheduler, QueryFrontend, Querier, Ingester, Distributor, Ruler, Compactor},

pkg/loki/modules.go

@@ -13,7 +13,6 @@ import (
 	"github.com/NYTimes/gziphandler"
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/grafana/dskit/kv"
 	"github.com/grafana/dskit/kv/codec"
 	"github.com/grafana/dskit/kv/memberlist"
 	"github.com/grafana/dskit/ring"
@@ -386,7 +385,9 @@ func (t *Loki) initStore() (_ services.Service, err error) {
 	// Always set these configs
 	// TODO(owen-d): Do the same for TSDB when we add IndexGatewayClientConfig
 	t.Cfg.StorageConfig.BoltDBShipperConfig.IndexGatewayClientConfig.Mode = t.Cfg.IndexGateway.Mode
-	t.Cfg.StorageConfig.BoltDBShipperConfig.IndexGatewayClientConfig.Ring = t.indexGatewayRing
+	if t.Cfg.IndexGateway.Mode == indexgateway.RingMode {
+		t.Cfg.StorageConfig.BoltDBShipperConfig.IndexGatewayClientConfig.Ring = t.indexGatewayRingManager.Ring
+	}
 
 	// If RF > 1 and current or upcoming index type is boltdb-shipper then disable index dedupe and write dedupe cache.
 	// This is to ensure that index entries are replicated to all the boltdb files in ingesters flushing replicated data.
@@ -855,7 +856,7 @@ func (t *Loki) initIndexGateway() (services.Service, error) {
 	t.Cfg.IndexGateway.Ring.KVStore.MemberlistKV = t.MemberlistKV.GetMemberlistKV
 	t.Cfg.IndexGateway.Ring.ListenPort = t.Cfg.Server.GRPCListenPort
 
-	indexClient, err := storage.NewIndexClient(config.BoltDBShipperType, t.Cfg.StorageConfig, t.Cfg.SchemaConfig, t.overrides, t.clientMetrics, prometheus.DefaultRegisterer)
+	indexClient, err := storage.NewIndexClient(config.BoltDBShipperType, t.Cfg.StorageConfig, t.Cfg.SchemaConfig, t.overrides, t.clientMetrics, t.indexGatewayRingManager.IndexGatewayOwnsTenant, prometheus.DefaultRegisterer)
 	if err != nil {
 		return nil, err
 	}
@@ -864,7 +865,7 @@ func (t *Loki) initIndexGateway() (services.Service, error) {
 		return nil, err
 	}
 
-	t.Server.HTTP.Path("/indexgateway/ring").Methods("GET", "POST").Handler(gateway)
+	t.Server.HTTP.Path("/indexgateway/ring").Methods("GET", "POST").Handler(t.indexGatewayRingManager)
 
 	indexgatewaypb.RegisterIndexGatewayServer(t.Server.GRPC, gateway)
 	return gateway, nil
@@ -878,29 +879,21 @@ func (t *Loki) initIndexGatewayRing() (_ services.Service, err error) {
 	t.Cfg.StorageConfig.BoltDBShipperConfig.Mode = shipper.ModeReadOnly
 	t.Cfg.IndexGateway.Ring.KVStore.MemberlistKV = t.MemberlistKV.GetMemberlistKV
 	t.Cfg.IndexGateway.Ring.ListenPort = t.Cfg.Server.GRPCListenPort
-	ringCfg := t.Cfg.IndexGateway.Ring.ToRingConfig(t.Cfg.IndexGateway.Ring.ReplicationFactor)
-	reg := prometheus.WrapRegistererWithPrefix("loki_", prometheus.DefaultRegisterer)
 
-	logger := util_log.Logger
-	ringStore, err := kv.NewClient(
-		ringCfg.KVStore,
-		ring.GetCodec(),
-		kv.RegistererWithKVName(prometheus.WrapRegistererWithPrefix("loki_", reg), "index-gateway"),
-		logger,
-	)
-	if err != nil {
-		return nil, gerrors.Wrap(err, "kv new client")
+	managerMode := indexgateway.ClientMode
+	if t.Cfg.isModuleEnabled(IndexGateway) {
+		managerMode = indexgateway.ServerMode
 	}
+	rm, err := indexgateway.NewRingManager(managerMode, t.Cfg.IndexGateway, util_log.Logger, prometheus.DefaultRegisterer)
 
-	t.indexGatewayRing, err = ring.NewWithStoreClientAndStrategy(
-		ringCfg, indexgateway.RingIdentifier, indexgateway.RingKey, ringStore, ring.NewIgnoreUnhealthyInstancesReplicationStrategy(), prometheus.WrapRegistererWithPrefix("loki_", reg), logger,
-	)
 	if err != nil {
-		return nil, gerrors.Wrap(err, "new with store client and strategy")
+		return nil, gerrors.Wrap(err, "new index gateway ring manager")
 	}
 
-	t.Server.HTTP.Path("/indexgateway/ring").Methods("GET", "POST").Handler(t.indexGatewayRing)
-	return t.indexGatewayRing, nil
+	t.indexGatewayRingManager = rm
+
+	t.Server.HTTP.Path("/indexgateway/ring").Methods("GET", "POST").Handler(t.indexGatewayRingManager)
+	return t.indexGatewayRingManager, nil
 }
 
 func (t *Loki) initQueryScheduler() (services.Service, error) {
@@ -962,7 +955,7 @@ func (t *Loki) deleteRequestsStore() (deletion.DeleteRequestsStore, error) {
 
 	deleteStore := deletion.NewNoOpDeleteRequestsStore()
 	if config.UsingBoltdbShipper(t.Cfg.SchemaConfig.Configs) && filteringEnabled {
-		indexClient, err := storage.NewIndexClient(config.BoltDBShipperType, t.Cfg.StorageConfig, t.Cfg.SchemaConfig, t.overrides, t.clientMetrics, prometheus.DefaultRegisterer)
+		indexClient, err := storage.NewIndexClient(config.BoltDBShipperType, t.Cfg.StorageConfig, t.Cfg.SchemaConfig, t.overrides, t.clientMetrics, nil, prometheus.DefaultRegisterer)
 		if err != nil {
 			return nil, err
 		}

pkg/storage/factory.go

@@ -124,7 +124,7 @@ func (cfg *Config) Validate() error {
 }
 
 // NewIndexClient makes a new index client of the desired type.
-func NewIndexClient(name string, cfg Config, schemaCfg config.SchemaConfig, limits StoreLimits, cm ClientMetrics, registerer prometheus.Registerer) (index.Client, error) {
+func NewIndexClient(name string, cfg Config, schemaCfg config.SchemaConfig, limits StoreLimits, cm ClientMetrics, ownsTenantFn downloads.IndexGatewayOwnsTenant, registerer prometheus.Registerer) (index.Client, error) {
 	switch name {
 	case config.StorageTypeInMemory:
 		store := testutils.NewMockStorage()
@@ -171,7 +171,7 @@ func NewIndexClient(name string, cfg Config, schemaCfg config.SchemaConfig, limi
 			return nil, err
 		}
 
-		boltDBIndexClientWithShipper, err = shipper.NewShipper(cfg.BoltDBShipperConfig, objectClient, limits, registerer)
+		boltDBIndexClientWithShipper, err = shipper.NewShipper(cfg.BoltDBShipperConfig, objectClient, limits, ownsTenantFn, registerer)
 
 		return boltDBIndexClientWithShipper, err
 	default:

pkg/storage/store.go

@@ -255,7 +255,7 @@ func (s *store) storeForPeriod(p config.PeriodConfig, chunkClient client.Client,
 			}, nil
 	}
 
-	idx, err := NewIndexClient(p.IndexType, s.cfg, s.schemaCfg, s.limits, s.clientMetrics, indexClientReg)
+	idx, err := NewIndexClient(p.IndexType, s.cfg, s.schemaCfg, s.limits, s.clientMetrics, nil, indexClientReg)
 	if err != nil {
 		return nil, nil, nil, errors.Wrap(err, "error creating index client")
 	}

pkg/storage/stores/shipper/downloads/table_manager.go

@@ -7,7 +7,6 @@ import (
 	"path/filepath"
 	"regexp"
 	"strconv"
-	"strings"
 	"sync"
 	"time"
 
@@ -32,6 +31,11 @@ type Limits interface {
 	DefaultLimits() *validation.Limits
 }
 
+// IndexGatewayOwnsTenant is invoked by an IndexGateway instance and answers whether if the given tenant is assigned to this instance or not.
+//
+// It is only relevant by an IndexGateway in the ring mode and if it returns false for a given tenant, that tenant will be ignored by this IndexGateway.
+type IndexGatewayOwnsTenant func(tenant string) bool
+
 type Config struct {
 	CacheDir          string
 	SyncInterval      time.Duration
@@ -52,9 +56,11 @@ type TableManager struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 	wg     sync.WaitGroup
+
+	ownsTenant IndexGatewayOwnsTenant
 }
 
-func NewTableManager(cfg Config, boltIndexClient BoltDBIndexClient, indexStorageClient storage.Client, registerer prometheus.Registerer) (*TableManager, error) {
+func NewTableManager(cfg Config, boltIndexClient BoltDBIndexClient, indexStorageClient storage.Client, ownsTenantFn IndexGatewayOwnsTenant, registerer prometheus.Registerer) (*TableManager, error) {
 	if err := chunk_util.EnsureDirectory(cfg.CacheDir); err != nil {
 		return nil, err
 	}
@@ -64,6 +70,7 @@ func NewTableManager(cfg Config, boltIndexClient BoltDBIndexClient, indexStorage
 		cfg:                cfg,
 		boltIndexClient:    boltIndexClient,
 		indexStorageClient: indexStorageClient,
+		ownsTenant:         ownsTenantFn,
 		tables:             make(map[string]Table),
 		metrics:            newMetrics(registerer),
 		ctx:                ctx,
@@ -319,8 +326,7 @@ func (tm *TableManager) ensureQueryReadiness(ctx context.Context) error {
 		if err := table.EnsureQueryReadiness(ctx, usersToBeQueryReadyFor); err != nil {
 			return err
 		}
-		joinedUsers := strings.Join(usersToBeQueryReadyFor, ",")
-		level.Info(util_log.Logger).Log("msg", "index pre-download for query readiness completed", "users", joinedUsers, "duration", time.Since(perTableStart), "table", tableName)
+		level.Info(util_log.Logger).Log("msg", "index pre-download for query readiness completed", "users_len", len(usersToBeQueryReadyFor), "duration", time.Since(perTableStart), "table", tableName)
 	}
 
 	return nil
@@ -345,6 +351,10 @@ func (tm *TableManager) findUsersInTableForQueryReadiness(tableNumber int64, use
 			continue
 		}
 
+		if tm.ownsTenant != nil && !tm.ownsTenant(userID) {
+			continue
+		}
+
 		if activeTableNumber-tableNumber <= int64(queryReadyNumDays) {
 			usersToBeQueryReadyFor = append(usersToBeQueryReadyFor, userID)
 		}

pkg/storage/stores/shipper/downloads/table_manager_test.go

@@ -26,7 +26,7 @@ func buildTestTableManager(t *testing.T, path string) (*TableManager, stopFunc)
 		CacheTTL:     time.Hour,
 		Limits:       &mockLimits{},
 	}
-	tableManager, err := NewTableManager(cfg, boltDBIndexClient, indexStorageClient, nil)
+	tableManager, err := NewTableManager(cfg, boltDBIndexClient, indexStorageClient, nil, nil)
 	require.NoError(t, err)
 
 	return tableManager, func() {

pkg/storage/stores/shipper/gateway_client.go

@@ -67,6 +67,10 @@ type IndexGatewayClientConfig struct {
 	// Forcefully disable the use of the index gateway client for the storage.
 	// This is mainly useful for the index-gateway component which should always use the storage.
 	Disabled bool `yaml:"-"`
+
+	// LogGatewayRequests configures if requests sent to the gateway should be logged or not.
+	// The log messages are of type debug and contain the address of the gateway and the relevant tenant.
+	LogGatewayRequests bool `yaml:"log_gateway_requests"`
 }
 
 // RegisterFlagsWithPrefix register client-specific flags with the given prefix.
@@ -75,6 +79,7 @@ type IndexGatewayClientConfig struct {
 func (i *IndexGatewayClientConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
 	i.GRPCClientConfig.RegisterFlagsWithPrefix(prefix+".grpc", f)
 	f.StringVar(&i.Address, prefix+".server-address", "", "Hostname or IP of the Index Gateway gRPC server running in simple mode.")
+	f.BoolVar(&i.LogGatewayRequests, prefix+".log-gateway-requests", false, "Whether requests sent to the gateway should be logged or not.")
 }
 
 func (i *IndexGatewayClientConfig) RegisterFlags(f *flag.FlagSet) {
@@ -310,6 +315,10 @@ func (s *GatewayClient) ringModeDo(ctx context.Context, callback func(client ind
 			continue
 		}
 
+		if s.cfg.LogGatewayRequests {
+			level.Debug(util_log.Logger).Log("msg", "sending request to gateway", "gateway", addr, "tenant", userID)
+		}
+
 		return nil
 	}
 

pkg/storage/stores/shipper/gateway_client_test.go

@@ -233,7 +233,7 @@ func benchmarkIndexQueries(b *testing.B, queries []index.Query) {
 		CacheTTL:          15 * time.Minute,
 		QueryReadyNumDays: 30,
 		Limits:            mockLimits{},
-	}, bclient, storage.NewIndexStorageClient(fs, "index/"), nil)
+	}, bclient, storage.NewIndexStorageClient(fs, "index/"), nil, nil)
 	require.NoError(b, err)
 
 	// initialize the index gateway server