allow logical plan fragment type for scheduler queue

Signed-off-by: rubywtl <[email protected]>

Author: rubywtl

pkg/cortex/modules.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"flag"
 	"fmt"
-
 	"log/slog"
 	"net/http"
 	"runtime"
@@ -414,6 +413,9 @@ func (t *Cortex) initQuerier() (serv services.Service, err error) {
 
 	t.Cfg.Worker.MaxConcurrentRequests = t.Cfg.Querier.MaxConcurrent
 	t.Cfg.Worker.TargetHeaders = t.Cfg.API.HTTPRequestHeadersToLog
+
+	t.Cfg.Worker.QuerierRing.ListenPort = t.Cfg.Server.GRPCListenPort
+
 	return querier_worker.NewQuerierWorker(t.Cfg.Worker, httpgrpc_server.NewServer(internalQuerierRouter), util_log.Logger, prometheus.DefaultRegisterer)
 }
 
@@ -813,7 +815,7 @@ func (t *Cortex) initQueryScheduler() (services.Service, error) {
 		tenant.WithDefaultResolver(tenantfederation.NewRegexValidator())
 	}
 
-	s, err := scheduler.NewScheduler(t.Cfg.QueryScheduler, t.Overrides, util_log.Logger, prometheus.DefaultRegisterer)
+	s, err := scheduler.NewScheduler(t.Cfg.QueryScheduler, t.Overrides, util_log.Logger, prometheus.DefaultRegisterer, t.Cfg.Querier.DistributedExecEnabled)
 	if err != nil {
 		return nil, errors.Wrap(err, "query-scheduler init")
 	}

pkg/distributed_execution/fragment_key.go

@@ -0,0 +1,33 @@
+package distributed_execution
+
+// FragmentKey uniquely identifies a fragment of a distributed logical query plan.
+// It combines a queryID (to identify the overall query) and a fragmentID
+// (to identify the specific fragment within that query).
+type FragmentKey struct {
+	// QueryID  identifies the distributed query this fragment belongs to
+	queryID uint64
+
+	// FragmentID identifies this specific fragment within the query
+	fragmentID uint64
+}
+
+// MakeFragmentKey creates a new FragmentKey with the given queryID and fragmentID.
+// It's used to track and identify fragments during distributed query execution.
+func MakeFragmentKey(queryID uint64, fragmentID uint64) FragmentKey {
+	return FragmentKey{
+		queryID:    queryID,
+		fragmentID: fragmentID,
+	}
+}
+
+// GetQueryID returns the queryID for the current key
+// This ID is shared across all fragments of the same distributed query.
+func (f FragmentKey) GetQueryID() uint64 {
+	return f.queryID
+}
+
+// GetFragmentID returns the ID for this specific fragment
+// within its parent query.
+func (f FragmentKey) GetFragmentID() uint64 {
+	return f.fragmentID
+}

pkg/distributed_execution/plan_fragments/fragmenter.go

@@ -0,0 +1,49 @@
+package plan_fragments
+
+import "github.com/thanos-io/promql-engine/logicalplan"
+
+type Fragmenter interface {
+	Fragment(node logicalplan.Node) ([]Fragment, error)
+}
+
+type DummyFragmenter struct {
+}
+
+func (f *DummyFragmenter) Fragment(node logicalplan.Node) ([]Fragment, error) {
+	// simple logic without distributed optimizer
+	return []Fragment{
+		{
+			Node:       node,
+			FragmentID: uint64(1),
+			ChildIDs:   []uint64{},
+			IsRoot:     true,
+		},
+	}, nil
+}
+
+type Fragment struct {
+	Node       logicalplan.Node
+	FragmentID uint64
+	ChildIDs   []uint64
+	IsRoot     bool
+}
+
+func (s *Fragment) IsEmpty() bool {
+	if s.Node != nil {
+		return false
+	}
+	if s.FragmentID != 0 {
+		return false
+	}
+	if s.IsRoot {
+		return false
+	}
+	if len(s.ChildIDs) != 0 {
+		return false
+	}
+	return true
+}
+
+func NewDummyFragmenter() Fragmenter {
+	return &DummyFragmenter{}
+}

pkg/distributed_execution/plan_fragments/fragmenter_test.go

@@ -0,0 +1,46 @@
+package plan_fragments
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/cortexproject/cortex/pkg/util/logical_plan"
+)
+
+func TestFragmenter(t *testing.T) {
+	type testCase struct {
+		name              string
+		query             string
+		start             time.Time
+		end               time.Time
+		expectedFragments int
+	}
+
+	now := time.Now()
+
+	// more tests will be added when distributed optimizer and fragmenter are implemented
+	tests := []testCase{
+		{
+			name:              "simple logical query plan - no fragmentation",
+			query:             "up",
+			start:             now,
+			end:               now,
+			expectedFragments: 1,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			lp, err := logical_plan.CreateTestLogicalPlan(tc.query, tc.start, tc.end, 0)
+			require.NoError(t, err)
+
+			fragmenter := NewDummyFragmenter()
+			res, err := fragmenter.Fragment((*lp).Root())
+
+			require.NoError(t, err)
+			require.Equal(t, tc.expectedFragments, len(res))
+		})
+	}
+}

pkg/querier/querier.go

@@ -96,7 +96,8 @@ type Config struct {
 	ParquetQueryableShardCacheSize    int    `yaml:"parquet_queryable_shard_cache_size"`
 	ParquetQueryableDefaultBlockStore string `yaml:"parquet_queryable_default_block_store"`
 	ParquetQueryableFallbackDisabled  bool   `yaml:"parquet_queryable_fallback_disabled"`
-	DistributedExecEnabled            bool   `yaml:"distributed_exec_enabled" doc:"hidden"`
+
+	DistributedExecEnabled bool `yaml:"distributed_exec_enabled" doc:"hidden"`
 }
 
 var (

pkg/querier/worker/scheduler_processor.go

@@ -31,7 +31,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/util/services"
 )
 
-func newSchedulerProcessor(cfg Config, handler RequestHandler, log log.Logger, reg prometheus.Registerer) (*schedulerProcessor, []services.Service) {
+func newSchedulerProcessor(cfg Config, handler RequestHandler, log log.Logger, reg prometheus.Registerer, querierAddress string) (*schedulerProcessor, []services.Service) {
 	p := &schedulerProcessor{
 		log:            log,
 		handler:        handler,
@@ -47,6 +47,7 @@ func newSchedulerProcessor(cfg Config, handler RequestHandler, log log.Logger, r
 			Help:    "Time spend doing requests to frontend.",
 			Buckets: prometheus.ExponentialBuckets(0.001, 4, 6),
 		}, []string{"operation", "status_code"}),
+		querierAddress: querierAddress,
 	}
 
 	frontendClientsGauge := promauto.With(reg).NewGauge(prometheus.GaugeOpts{
@@ -71,6 +72,7 @@ type schedulerProcessor struct {
 	grpcConfig     grpcclient.Config
 	maxMessageSize int
 	querierID      string
+	querierAddress string
 
 	frontendPool                  *client.Pool
 	frontendClientRequestDuration *prometheus.HistogramVec
@@ -97,7 +99,7 @@ func (sp *schedulerProcessor) processQueriesOnSingleStream(ctx context.Context,
 	for backoff.Ongoing() {
 		c, err := schedulerClient.QuerierLoop(ctx)
 		if err == nil {
-			err = c.Send(&schedulerpb.QuerierToScheduler{QuerierID: sp.querierID})
+			err = c.Send(&schedulerpb.QuerierToScheduler{QuerierID: sp.querierID, QuerierAddress: sp.querierAddress})
 		}
 
 		if err != nil {

pkg/querier/worker/scheduler_processor_test.go

@@ -144,7 +144,7 @@ func Test_ToShowNotPanic_RelatedIssue6599(t *testing.T) {
 		go stat.AddFetchedChunkBytes(10)
 	}).Return(&httpgrpc.HTTPResponse{}, nil)
 
-	sp, _ := newSchedulerProcessor(cfg, requestHandler, log.NewNopLogger(), nil)
+	sp, _ := newSchedulerProcessor(cfg, requestHandler, log.NewNopLogger(), nil, "")
 	schedulerClient := &mockSchedulerForQuerierClient{}
 	schedulerClient.On("QuerierLoop", mock.Anything, mock.Anything).Return(querierLoopClient, nil)
 

pkg/querier/worker/worker.go

@@ -3,7 +3,10 @@ package worker
 import (
 	"context"
 	"flag"
+	"fmt"
+	"net"
 	"os"
+	"strconv"
 	"sync"
 	"time"
 
@@ -14,7 +17,10 @@ import (
 	"github.com/weaveworks/common/httpgrpc"
 	"google.golang.org/grpc"
 
+	"github.com/cortexproject/cortex/pkg/ring"
+	"github.com/cortexproject/cortex/pkg/ring/kv"
 	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/pkg/util/flagext"
 	"github.com/cortexproject/cortex/pkg/util/grpcclient"
 	"github.com/cortexproject/cortex/pkg/util/services"
 )
@@ -33,9 +39,71 @@ type Config struct {
 	GRPCClientConfig grpcclient.Config `yaml:"grpc_client_config"`
 
 	TargetHeaders []string `yaml:"-"` // Propagated by config.
+
+	QuerierRing RingConfig `yaml:"ring"`
+}
+
+type RingConfig struct {
+	KVStore                kv.Config     `yaml:"kvstore" doc:"description=The key-value store used to share the hash ring across multiple instances."`
+	HeartbeatPeriod        time.Duration `yaml:"heartbeat_period"`
+	HeartbeatTimeout       time.Duration `yaml:"heartbeat_timeout"`
+	ReplicationFactor      int           `yaml:"replication_factor"`
+	ZoneAwarenessEnabled   bool          `yaml:"zone_awareness_enabled"`
+	TokensFilePath         string        `yaml:"tokens_file_path"`
+	DetailedMetricsEnabled bool          `yaml:"detailed_metrics_enabled"`
+
+	FinalSleep                      time.Duration `yaml:"final_sleep"`
+	WaitInstanceStateTimeout        time.Duration `yaml:"wait_instance_state_timeout"`
+	KeepInstanceInTheRingOnShutdown bool          `yaml:"keep_instance_in_the_ring_on_shutdown"`
+
+	// Instance details
+	InstanceID             string   `yaml:"instance_id" doc:"hidden"`
+	InstanceInterfaceNames []string `yaml:"instance_interface_names"`
+	InstancePort           int      `yaml:"instance_port" doc:"hidden"`
+	InstanceAddr           string   `yaml:"instance_addr" doc:"hidden"`
+	InstanceZone           string   `yaml:"instance_availability_zone"`
+
+	NumTokens int `yaml:"num_tokens"`
+
+	// Injected internally
+	ListenPort      int           `yaml:"-"`
+	RingCheckPeriod time.Duration `yaml:"-"`
+
+	// Used for testing
+	SkipUnregister bool `yaml:"-"`
+}
+
+// RegisterFlags adds the flags required to config this to the given FlagSet
+func (cfg *RingConfig) RegisterFlags(f *flag.FlagSet) {
+	hostname, err := os.Hostname()
+	if err != nil {
+		panic(fmt.Errorf("failed to get hostname, %w", err))
+	}
+
+	// Ring flags
+	cfg.KVStore.RegisterFlagsWithPrefix("querier.ring.", "querier/", f)
+	f.DurationVar(&cfg.HeartbeatPeriod, "querier.ring.heartbeat-period", 5*time.Second, "Period at which to heartbeat to the ring. 0 = disabled.")
+	f.DurationVar(&cfg.HeartbeatTimeout, "querier.ring.heartbeat-timeout", time.Minute, "The heartbeat timeout after which rulers are considered unhealthy within the ring. 0 = never (timeout disabled).")
+	f.DurationVar(&cfg.FinalSleep, "querier.ring.final-sleep", 0*time.Second, "The sleep seconds when ruler is shutting down. Need to be close to or larger than KV Store information propagation delay")
+	f.IntVar(&cfg.ReplicationFactor, "querier.ring.replication-factor", 1, "EXPERIMENTAL: The replication factor to use when loading rule groups for API HA.")
+	f.BoolVar(&cfg.ZoneAwarenessEnabled, "querier.ring.zone-awareness-enabled", false, "EXPERIMENTAL: True to enable zone-awareness and load rule groups across different availability zones for API HA.")
+	f.StringVar(&cfg.TokensFilePath, "querier.ring.tokens-file-path", "", "EXPERIMENTAL: File path where tokens are stored. If empty, tokens are not stored at shutdown and restored at startup.")
+	f.BoolVar(&cfg.DetailedMetricsEnabled, "querier.ring.detailed-metrics-enabled", true, "Set to true to enable ring detailed metrics. These metrics provide detailed information, such as token count and ownership per tenant. Disabling them can significantly decrease the number of metrics emitted.")
+
+	// Instance flags
+	cfg.InstanceInterfaceNames = []string{"eth0", "en0"}
+	f.Var((*flagext.StringSlice)(&cfg.InstanceInterfaceNames), "querier.ring.instance-interface-names", "Name of network interface to read address from.")
+	f.StringVar(&cfg.InstanceAddr, "querier.ring.instance-addr", "", "IP address to advertise in the ring.")
+	f.IntVar(&cfg.InstancePort, "querier.ring.instance-port", 0, "Port to advertise in the ring (defaults to server.grpc-listen-port).")
+	f.StringVar(&cfg.InstanceID, "querier.ring.instance-id", hostname, "Instance ID to register in the ring.")
+	f.StringVar(&cfg.InstanceZone, "querier.ring.instance-availability-zone", "", "The availability zone where this instance is running. Required if zone-awareness is enabled.")
+	f.IntVar(&cfg.NumTokens, "querier.ring.num-tokens", 128, "Number of tokens for each ruler.")
+	f.BoolVar(&cfg.KeepInstanceInTheRingOnShutdown, "querier.ring.keep-instance-in-the-ring-on-shutdown", false, "Keep instance in the ring on shut down.")
 }
 
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
+	cfg.QuerierRing.RegisterFlags(f)
+
 	f.StringVar(&cfg.SchedulerAddress, "querier.scheduler-address", "", "Hostname (and port) of scheduler that querier will periodically resolve, connect to and receive queries from. Only one of -querier.frontend-address or -querier.scheduler-address can be set. If neither is set, queries are only received via HTTP endpoint.")
 	f.StringVar(&cfg.FrontendAddress, "querier.frontend-address", "", "Address of query frontend service, in host:port format. If -querier.scheduler-address is set as well, querier will use scheduler instead. Only one of -querier.frontend-address or -querier.scheduler-address can be set. If neither is set, queries are only received via HTTP endpoint.")
 
@@ -109,7 +177,14 @@ func NewQuerierWorker(cfg Config, handler RequestHandler, log log.Logger, reg pr
 		level.Info(log).Log("msg", "Starting querier worker connected to query-scheduler", "scheduler", cfg.SchedulerAddress)
 
 		address = cfg.SchedulerAddress
-		processor, servs = newSchedulerProcessor(cfg, handler, log, reg)
+
+		ipAddr, err := ring.GetInstanceAddr(cfg.QuerierRing.InstanceAddr, cfg.QuerierRing.InstanceInterfaceNames, log)
+		if err != nil {
+			return nil, err
+		}
+		querierAddr := net.JoinHostPort(ipAddr, strconv.Itoa(cfg.QuerierRing.ListenPort))
+
+		processor, servs = newSchedulerProcessor(cfg, handler, log, reg, querierAddr)
 
 	case cfg.FrontendAddress != "":
 		level.Info(log).Log("msg", "Starting querier worker connected to query-frontend", "frontend", cfg.FrontendAddress)

pkg/scheduler/fragment_table/fragment_table.go

@@ -0,0 +1,70 @@
+package fragment_table
+
+import (
+	"sync"
+
+	"github.com/cortexproject/cortex/pkg/distributed_execution"
+)
+
+type FragmentTable struct {
+	mappings map[distributed_execution.FragmentKey]string
+	mu       sync.RWMutex
+}
+
+func NewFragmentTable() *FragmentTable {
+	return &FragmentTable{
+		mappings: make(map[distributed_execution.FragmentKey]string),
+	}
+}
+
+func (f *FragmentTable) AddMapping(queryID uint64, fragmentID uint64, addr string) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+	f.mappings[key] = addr
+}
+
+func (f *FragmentTable) GetAllChildAddresses(queryID uint64, fragmentIDs []uint64) ([]string, bool) {
+	f.mu.RLock()
+	defer f.mu.RUnlock()
+
+	addresses := make([]string, 0, len(fragmentIDs))
+
+	for _, fragmentID := range fragmentIDs {
+		key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+		if addr, ok := f.mappings[key]; ok {
+			addresses = append(addresses, addr)
+		} else {
+			return nil, false
+		}
+	}
+	return addresses, true
+}
+
+func (f *FragmentTable) GetChildAddr(queryID uint64, fragmentID uint64) (string, bool) {
+	f.mu.RLock()
+	defer f.mu.RUnlock()
+
+	key := distributed_execution.MakeFragmentKey(queryID, fragmentID)
+	if addr, ok := f.mappings[key]; ok {
+		return addr, true
+	}
+	return "", false
+}
+
+func (f *FragmentTable) ClearMappings(queryID uint64) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	keysToDelete := make([]distributed_execution.FragmentKey, 0)
+	for key := range f.mappings {
+		if key.GetQueryID() == queryID {
+			keysToDelete = append(keysToDelete, key)
+		}
+	}
+
+	for _, key := range keysToDelete {
+		delete(f.mappings, key)
+	}
+}