feat(scheduler): implement and register block builder rpc service

pkg/blockbuilder/builder/builder.go

@@ -297,7 +297,7 @@ func (i *BlockBuilder) processJob(ctx context.Context, job *types.Job, logger lo
 		"load records",
 		1,
 		func(ctx context.Context) error {
-			lastOffset, err = i.loadRecords(ctx, int32(job.Partition), job.Offsets, inputCh)
+			lastOffset, err = i.loadRecords(ctx, job.Partition, job.Offsets, inputCh)
 			return err
 		},
 		func(ctx context.Context) error {
@@ -545,7 +545,7 @@ func (i *BlockBuilder) loadRecords(ctx context.Context, partitionID int32, offse
 		}
 	}
 
-	return lastOffset, err
+	return lastOffset, boff.Err()
 }
 
 func withBackoff[T any](

pkg/blockbuilder/scheduler/queue.go

@@ -1,11 +1,13 @@
 package scheduler
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"time"
 
 	"github.com/grafana/loki/v3/pkg/blockbuilder/types"
+	"github.com/grafana/loki/v3/pkg/queue"
 )
 
 const (
@@ -44,18 +46,24 @@ type JobQueue struct {
 	completed  *CircularBuffer[*types.Job]           // Last N completed jobs
 	statusMap  map[string]types.JobStatus            // Maps job ID to its current status
 	mu         sync.RWMutex
+	cond       queue.ContextCond // to signal when a job is available
 }
 
 // NewJobQueue creates a new job queue instance
 func NewJobQueue() *JobQueue {
-	return &JobQueue{
+	q := &JobQueue{
 		pending: NewPriorityQueue[*JobWithPriority[int]](func(a, b *JobWithPriority[int]) bool {
 			return a.Priority > b.Priority // Higher priority first
 		}),
 		inProgress: make(map[string]*inProgressJob),
 		completed:  NewCircularBuffer[*types.Job](defaultCompletedJobsCapacity),
 		statusMap:  make(map[string]types.JobStatus),
 	}
+	q.cond = queue.ContextCond{
+		Cond: sync.NewCond(&q.mu),
+	}
+
+	return q
 }
 
 func (q *JobQueue) Exists(job *types.Job) (types.JobStatus, bool) {
@@ -79,15 +87,24 @@ func (q *JobQueue) Enqueue(job *types.Job, priority int) error {
 	jobWithPriority := NewJobWithPriority(job, priority)
 	q.pending.Push(jobWithPriority)
 	q.statusMap[job.ID] = types.JobStatusPending
+	q.cond.Broadcast()
 	return nil
 }
 
 // Dequeue gets the next available job and assigns it to a builder
-func (q *JobQueue) Dequeue(_ string) (*types.Job, bool, error) {
+func (q *JobQueue) Dequeue(ctx context.Context, _ string, wait bool) (*types.Job, bool, error) {
 	q.mu.Lock()
 	defer q.mu.Unlock()
 
-	if q.pending.Len() == 0 {
+	if wait {
+		for q.pending.Len() == 0 && ctx.Err() == nil {
+			q.cond.Wait(ctx)
+		}
+
+		if ctx.Err() != nil {
+			return nil, false, ctx.Err()
+		}
+	} else if q.pending.Len() == 0 {
 		return nil, false, nil
 	}
 

pkg/blockbuilder/scheduler/scheduler.go

@@ -14,10 +14,10 @@ import (
 	"github.com/twmb/franz-go/pkg/kadm"
 
 	"github.com/grafana/loki/v3/pkg/blockbuilder/types"
+	"github.com/grafana/loki/v3/pkg/blockbuilder/types/proto"
 )
 
 var (
-	_ types.Scheduler = unimplementedScheduler{}
 	_ types.Scheduler = &BlockScheduler{}
 )
 
@@ -140,7 +140,7 @@ func (s *BlockScheduler) HandleGetJob(ctx context.Context, builderID string) (*t
 	case <-ctx.Done():
 		return nil, false, ctx.Err()
 	default:
-		return s.queue.Dequeue(builderID)
+		return s.queue.Dequeue(ctx, builderID, false)
 	}
 }
 
@@ -155,17 +155,44 @@ func (s *BlockScheduler) HandleSyncJob(_ context.Context, builderID string, job
 	return nil
 }
 
-// unimplementedScheduler provides default implementations that panic.
-type unimplementedScheduler struct{}
-
-func (s unimplementedScheduler) HandleGetJob(_ context.Context, _ string) (*types.Job, bool, error) {
-	panic("unimplemented")
+func (s *BlockScheduler) CompleteJob(_ context.Context, req *proto.CompleteJobRequest) (*proto.CompleteJobResponse, error) {
+	s.queue.MarkComplete(req.Job.Id)
+	return nil, nil
 }
 
-func (s unimplementedScheduler) HandleCompleteJob(_ context.Context, _ string, _ *types.Job) error {
-	panic("unimplemented")
+func (s *BlockScheduler) SyncJob(_ context.Context, req *proto.SyncJobRequest) (*proto.SyncJobResponse, error) {
+	s.queue.SyncJob(req.Job.Id, req.BuilderId, &types.Job{
+		ID:        req.Job.Id,
+		Partition: req.Job.Partition,
+		Offsets: types.Offsets{
+			Min: req.Job.Offsets.Min,
+			Max: req.Job.Offsets.Max,
+		},
+	})
+
+	return nil, nil
 }
 
-func (s unimplementedScheduler) HandleSyncJob(_ context.Context, _ string, _ *types.Job) error {
-	panic("unimplemented")
+func (s *BlockScheduler) GetJob(req *proto.GetJobRequest, stream proto.BlockBuilderService_GetJobServer) error {
+	job, ok, err := s.queue.Dequeue(stream.Context(), req.BuilderId, true)
+	if err != nil {
+		return err
+	}
+
+	var jobpb *proto.Job
+	if ok {
+		jobpb = &proto.Job{
+			Id:        job.ID,
+			Partition: job.Partition,
+			Offsets: &proto.Offsets{
+				Min: job.Offsets.Min,
+				Max: job.Offsets.Max,
+			},
+		}
+	}
+
+	return stream.Send(&proto.GetJobResponse{
+		Job: jobpb,
+		Ok:  ok,
+	})
 }

pkg/blockbuilder/scheduler/strategy.go

@@ -60,7 +60,7 @@ func (p *RecordCountPlanner) Plan(ctx context.Context) ([]*JobWithPriority[int],
 		endOffset := min(startOffset+p.targetRecordCount, partitionOffset.End.Offset)
 
 		job := NewJobWithPriority(
-			types.NewJob(int(partitionOffset.Partition), types.Offsets{
+			types.NewJob(partitionOffset.Partition, types.Offsets{
 				Min: startOffset,
 				Max: endOffset,
 			}), int(partitionOffset.End.Offset-startOffset),
@@ -135,7 +135,7 @@ func (p *TimeRangePlanner) Plan(ctx context.Context) ([]*JobWithPriority[int], e
 		}
 
 		job := NewJobWithPriority(
-			types.NewJob(int(partitionOffset.Partition), types.Offsets{
+			types.NewJob(partitionOffset.Partition, types.Offsets{
 				Min: startOffset,
 				Max: endOffset,
 			}), int(endOffset-startOffset),

pkg/blockbuilder/types/grpc_transport.go

@@ -70,7 +70,12 @@ func (t *GRPCTransport) SendGetJobRequest(ctx context.Context, req *GetJobReques
 		BuilderId: req.BuilderID,
 	}
 
-	resp, err := t.GetJob(ctx, protoReq)
+	client, err := t.GetJob(ctx, protoReq)
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := client.Recv()
 	if err != nil {
 		return nil, err
 	}
@@ -111,7 +116,7 @@ func protoToJob(p *proto.Job) *Job {
 	}
 	return &Job{
 		ID:        p.GetId(),
-		Partition: int(p.GetPartition()),
+		Partition: p.GetPartition(),
 		Offsets: Offsets{
 			Min: p.GetOffsets().GetMin(),
 			Max: p.GetOffsets().GetMax(),
@@ -126,7 +131,7 @@ func jobToProto(j *Job) *proto.Job {
 	}
 	return &proto.Job{
 		Id:        j.ID,
-		Partition: int32(j.Partition),
+		Partition: j.Partition,
 		Offsets: &proto.Offsets{
 			Min: j.Offsets.Min,
 			Max: j.Offsets.Max,

pkg/blockbuilder/types/job.go

@@ -6,7 +6,7 @@ import "fmt"
 type Job struct {
 	ID string
 	// Partition and offset information
-	Partition int
+	Partition int32
 	Offsets   Offsets
 }
 
@@ -26,7 +26,7 @@ type Offsets struct {
 }
 
 // NewJob creates a new job with the given partition and offsets
-func NewJob(partition int, offsets Offsets) *Job {
+func NewJob(partition int32, offsets Offsets) *Job {
 	return &Job{
 		ID:        GenerateJobID(partition, offsets),
 		Partition: partition,
@@ -35,6 +35,6 @@ func NewJob(partition int, offsets Offsets) *Job {
 }
 
 // GenerateJobID creates a deterministic job ID from partition and offsets
-func GenerateJobID(partition int, offsets Offsets) string {
+func GenerateJobID(partition int32, offsets Offsets) string {
 	return fmt.Sprintf("job-%d-%d-%d", partition, offsets.Min, offsets.Max)
 }