allow logical plan fragment type for scheduler queue

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

Author: rubywtl

pkg/cortex/modules.go

@@ -6,9 +6,11 @@ import (
 	"fmt"
 
 	"log/slog"
+	"net"
 	"net/http"
 	"runtime"
 	"runtime/debug"
+	"strconv"
 
 	"github.com/go-kit/log/level"
 	"github.com/opentracing-contrib/go-stdlib/nethttp"
@@ -414,7 +416,12 @@ 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
-	return querier_worker.NewQuerierWorker(t.Cfg.Worker, httpgrpc_server.NewServer(internalQuerierRouter), util_log.Logger, prometheus.DefaultRegisterer)
+	ipAddr, err := ring.GetInstanceAddr(t.Cfg.Alertmanager.ShardingRing.InstanceAddr, t.Cfg.Alertmanager.ShardingRing.InstanceInterfaceNames, util_log.Logger)
+	if err != nil {
+		return nil, err
+	}
+	serverAddress := net.JoinHostPort(ipAddr, strconv.Itoa(t.Cfg.Server.GRPCListenPort))
+	return querier_worker.NewQuerierWorker(t.Cfg.Worker, httpgrpc_server.NewServer(internalQuerierRouter), util_log.Logger, prometheus.DefaultRegisterer, serverAddress)
 }
 
 func (t *Cortex) initStoreQueryables() (services.Service, error) {
@@ -813,7 +820,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,21 @@
+package distributed_execution
+
+type FragmentKey struct {
+	queryID    uint64
+	fragmentID uint64
+}
+
+func MakeFragmentKey(queryID uint64, fragmentID uint64) *FragmentKey {
+	return &FragmentKey{
+		queryID:    queryID,
+		fragmentID: fragmentID,
+	}
+}
+
+func (f FragmentKey) GetQueryID() uint64 {
+	return f.queryID
+}
+
+func (f FragmentKey) GetFragmentID() uint64 {
+	return f.fragmentID
+}

pkg/frontend/frontend_test.go

@@ -291,7 +291,7 @@ func testFrontend(t *testing.T, config CombinedFrontendConfig, handler http.Hand
 	go grpcServer.Serve(grpcListen) //nolint:errcheck
 
 	var worker services.Service
-	worker, err = querier_worker.NewQuerierWorker(workerConfig, httpgrpc_server.NewServer(handler), logger, nil)
+	worker, err = querier_worker.NewQuerierWorker(workerConfig, httpgrpc_server.NewServer(handler), logger, nil, "")
 	require.NoError(t, err)
 	require.NoError(t, services.StartAndAwaitRunning(context.Background(), worker))
 

pkg/frontend/v1/frontend_test.go

@@ -285,7 +285,7 @@ func testFrontend(t *testing.T, config Config, handler http.Handler, test func(a
 	go grpcServer.Serve(grpcListen) //nolint:errcheck
 
 	var worker services.Service
-	worker, err = querier_worker.NewQuerierWorker(workerConfig, httpgrpc_server.NewServer(handler), logger, nil)
+	worker, err = querier_worker.NewQuerierWorker(workerConfig, httpgrpc_server.NewServer(handler), logger, nil, "")
 	require.NoError(t, err)
 	require.NoError(t, services.StartAndAwaitRunning(context.Background(), worker))
 

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

@@ -91,7 +91,7 @@ type querierWorker struct {
 	managers map[string]*processorManager
 }
 
-func NewQuerierWorker(cfg Config, handler RequestHandler, log log.Logger, reg prometheus.Registerer) (services.Service, error) {
+func NewQuerierWorker(cfg Config, handler RequestHandler, log log.Logger, reg prometheus.Registerer, querierAddr string) (services.Service, error) {
 	if cfg.QuerierID == "" {
 		hostname, err := os.Hostname()
 		if err != nil {
@@ -109,7 +109,7 @@ 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)
+		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/plan_fragments/fragment_table.go

@@ -0,0 +1,70 @@
+package plan_fragments
+
+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) GetMappings(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) GetMapping(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)
+	}
+}

pkg/scheduler/plan_fragments/fragment_table_test.go

@@ -0,0 +1,116 @@
+package plan_fragments
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// TestSchedulerCoordination checks whether the hashtable for fragment-querier mapping gives the expected value
+// It also checks if it remains functional and accurate during a multi-thread/concurrent read & write situation
+func TestSchedulerCoordination(t *testing.T) {
+	t.Run("basic operations", func(t *testing.T) {
+		table := NewFragmentTable()
+		table.AddMapping(uint64(0), uint64(1), "localhost:8000")
+		table.AddMapping(uint64(0), uint64(2), "localhost:8001")
+
+		result, exist := table.GetMappings(uint64(0), []uint64{1, 2})
+		require.True(t, exist)
+		require.Equal(t, []string{"localhost:8000", "localhost:8001"}, result)
+
+		result, exist = table.GetMappings(uint64(0), []uint64{1, 3})
+		require.False(t, exist)
+		require.Empty(t, result)
+
+		result, exist = table.GetMappings(uint64(0), []uint64{1})
+		require.True(t, exist)
+		require.Equal(t, []string{"localhost:8000"}, result)
+
+		table.ClearMappings(uint64(0))
+		result, exist = table.GetMappings(uint64(0), []uint64{1})
+		require.False(t, exist)
+		require.Empty(t, result)
+	})
+
+	t.Run("concurrent operations", func(t *testing.T) {
+		table := NewFragmentTable()
+		const numGoroutines = 10
+		const numOperations = 100
+
+		var wg sync.WaitGroup
+		wg.Add(numGoroutines * 3)
+
+		// write
+		for i := 0; i < numGoroutines; i++ {
+			go func(routine int) {
+				defer wg.Done()
+				for j := 0; j < numOperations; j++ {
+					queryID := uint64(routine)
+					fragmentID := uint64(j)
+					addr := fmt.Sprintf("localhost:%d", j)
+					table.AddMapping(queryID, fragmentID, addr)
+				}
+			}(i)
+		}
+
+		// read
+		for i := 0; i < numGoroutines; i++ {
+			go func(routine int) {
+				defer wg.Done()
+				for j := 0; j < numOperations; j++ {
+					queryID := uint64(routine)
+					fragmentIDs := []uint64{uint64(j)}
+					table.GetMappings(queryID, fragmentIDs)
+				}
+			}(i)
+		}
+
+		// clear
+		for i := 0; i < numGoroutines; i++ {
+			go func(routine int) {
+				defer wg.Done()
+				for j := 0; j < numOperations; j++ {
+					queryID := uint64(routine)
+					table.ClearMappings(queryID)
+				}
+			}(i)
+		}
+
+		wg.Wait()
+	})
+
+	t.Run("edge cases", func(t *testing.T) {
+		table := NewFragmentTable()
+
+		// test empty fragment IDs
+		result, exist := table.GetMappings(0, []uint64{})
+		require.True(t, exist)
+		require.Empty(t, result)
+
+		// test clearing non-existent query
+		table.ClearMappings(999)
+		require.NotPanics(t, func() {
+			table.ClearMappings(999)
+		})
+
+		// test overwriting mapping
+		table.AddMapping(1, 1, "addr1")
+		table.AddMapping(1, 1, "addr2")
+		result, exist = table.GetMappings(1, []uint64{1})
+		require.True(t, exist)
+		require.Equal(t, []string{"addr2"}, result)
+
+		// test multiple queries
+		table.AddMapping(1, 1, "addr1")
+		table.AddMapping(2, 1, "addr2")
+		result, exist = table.GetMappings(1, []uint64{1})
+		require.True(t, exist)
+		require.Equal(t, []string{"addr1"}, result)
+
+		result, exist = table.GetMappings(2, []uint64{1})
+		require.True(t, exist)
+		require.Equal(t, []string{"addr2"}, result)
+	})
+}

pkg/scheduler/plan_fragments/fragmenter.go

@@ -0,0 +1,54 @@
+package plan_fragments
+
+import "github.com/thanos-io/promql-engine/logicalplan"
+
+type Fragmenter interface {
+	Fragment(node logicalplan.Node) ([]Fragment, error)
+	getNewID() uint64
+}
+
+type DummyFragmenter struct {
+}
+
+func (f *DummyFragmenter) Fragment(node logicalplan.Node) ([]Fragment, error) {
+	// simple logic without distributed optimizer
+	return []Fragment{
+		{
+			Node:       node,
+			FragmentID: f.getNewID(),
+			ChildIDs:   []uint64{},
+			IsRoot:     true,
+		},
+	}, nil
+}
+
+func (f *DummyFragmenter) getNewID() uint64 {
+	return 1 // for dummy plan_fragments testing
+}
+
+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{}
+}