feat(blockbuilder): grpc transport (#15218)

pkg/blockbuilder/scheduler/scheduler_test.go

@@ -15,15 +15,15 @@ import (
 type testEnv struct {
 	queue     *JobQueue
 	scheduler *BlockScheduler
-	transport *builder.MemoryTransport
+	transport *types.MemoryTransport
 	builder   *builder.Worker
 }
 
 func newTestEnv(builderID string) *testEnv {
 	queue := NewJobQueue()
 	scheduler := NewScheduler(Config{}, queue, nil, log.NewNopLogger(), prometheus.NewRegistry())
-	transport := builder.NewMemoryTransport(scheduler)
-	builder := builder.NewWorker(builderID, builder.NewMemoryTransport(scheduler))
+	transport := types.NewMemoryTransport(scheduler)
+	builder := builder.NewWorker(builderID, transport)
 
 	return &testEnv{
 		queue:     queue,
@@ -89,7 +89,7 @@ func TestMultipleBuilders(t *testing.T) {
 	// Create first environment
 	env1 := newTestEnv("test-builder-1")
 	// Create second builder using same scheduler
-	builder2 := builder.NewWorker("test-builder-2", builder.NewMemoryTransport(env1.scheduler))
+	builder2 := builder.NewWorker("test-builder-2", env1.transport)
 
 	ctx := context.Background()
 

pkg/blockbuilder/types/grpc_transport.go

@@ -0,0 +1,147 @@
+package types
+
+import (
+	"context"
+	"flag"
+	"io"
+
+	"github.com/grafana/dskit/grpcclient"
+	"github.com/grafana/dskit/instrument"
+	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/health/grpc_health_v1"
+
+	"github.com/grafana/loki/v3/pkg/blockbuilder/types/proto"
+	"github.com/grafana/loki/v3/pkg/util/constants"
+)
+
+var _ Transport = &GRPCTransport{}
+
+type GRPCTransportConfig struct {
+	Address string `yaml:"address,omitempty"`
+
+	// GRPCClientConfig configures the gRPC connection between the Bloom Gateway client and the server.
+	GRPCClientConfig grpcclient.Config `yaml:"grpc_client_config"`
+}
+
+func (cfg *GRPCTransportConfig) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
+	f.StringVar(&cfg.Address, prefix+"address", "", "address in DNS Service Discovery format: https://grafana.com/docs/mimir/latest/configure/about-dns-service-discovery/#supported-discovery-modes")
+}
+
+type grpcTransportMetrics struct {
+	requestLatency *prometheus.HistogramVec
+}
+
+func newGRPCTransportMetrics(registerer prometheus.Registerer) *grpcTransportMetrics {
+	return &grpcTransportMetrics{
+		requestLatency: promauto.With(registerer).NewHistogramVec(prometheus.HistogramOpts{
+			Namespace: constants.Loki,
+			Subsystem: "block_builder_grpc",
+			Name:      "request_duration_seconds",
+			Help:      "Time (in seconds) spent serving requests when using the block builder grpc transport",
+			Buckets:   instrument.DefBuckets,
+		}, []string{"operation", "status_code"}),
+	}
+}
+
+// GRPCTransport implements the Transport interface using gRPC
+type GRPCTransport struct {
+	grpc_health_v1.HealthClient
+	io.Closer
+	proto.BlockBuilderServiceClient
+}
+
+// NewGRPCTransportFromAddress creates a new gRPC transport instance from an address and dial options
+func NewGRPCTransportFromAddress(
+	metrics *grpcTransportMetrics,
+	cfg GRPCTransportConfig,
+) (*GRPCTransport, error) {
+
+	dialOpts, err := cfg.GRPCClientConfig.DialOption(grpcclient.Instrument(metrics.requestLatency))
+	if err != nil {
+		return nil, err
+	}
+
+	// nolint:staticcheck // grpc.Dial() has been deprecated; we'll address it before upgrading to gRPC 2.
+	conn, err := grpc.Dial(cfg.Address, dialOpts...)
+	if err != nil {
+		return nil, errors.Wrap(err, "new grpc pool dial")
+	}
+
+	return &GRPCTransport{
+		Closer:                    conn,
+		HealthClient:              grpc_health_v1.NewHealthClient(conn),
+		BlockBuilderServiceClient: proto.NewBlockBuilderServiceClient(conn),
+	}, nil
+}
+
+// SendGetJobRequest implements Transport
+func (t *GRPCTransport) SendGetJobRequest(ctx context.Context, req *GetJobRequest) (*GetJobResponse, error) {
+	protoReq := &proto.GetJobRequest{
+		BuilderId: req.BuilderID,
+	}
+
+	resp, err := t.GetJob(ctx, protoReq)
+	if err != nil {
+		return nil, err
+	}
+
+	return &GetJobResponse{
+		Job: protoToJob(resp.GetJob()),
+		OK:  resp.GetOk(),
+	}, nil
+}
+
+// SendCompleteJob implements Transport
+func (t *GRPCTransport) SendCompleteJob(ctx context.Context, req *CompleteJobRequest) error {
+	protoReq := &proto.CompleteJobRequest{
+		BuilderId: req.BuilderID,
+		Job:       jobToProto(req.Job),
+	}
+
+	_, err := t.CompleteJob(ctx, protoReq)
+	return err
+}
+
+// SendSyncJob implements Transport
+func (t *GRPCTransport) SendSyncJob(ctx context.Context, req *SyncJobRequest) error {
+	protoReq := &proto.SyncJobRequest{
+		BuilderId: req.BuilderID,
+		Job:       jobToProto(req.Job),
+	}
+
+	_, err := t.SyncJob(ctx, protoReq)
+	return err
+}
+
+// protoToJob converts a proto Job to a types.Job
+func protoToJob(p *proto.Job) *Job {
+	if p == nil {
+		return nil
+	}
+	return &Job{
+		ID:        p.GetId(),
+		Partition: int(p.GetPartition()),
+		Offsets: Offsets{
+			Min: p.GetOffsets().GetMin(),
+			Max: p.GetOffsets().GetMax(),
+		},
+	}
+}
+
+// jobToProto converts a types.Job to a proto Job
+func jobToProto(j *Job) *proto.Job {
+	if j == nil {
+		return nil
+	}
+	return &proto.Job{
+		Id:        j.ID,
+		Partition: int32(j.Partition),
+		Offsets: &proto.Offsets{
+			Min: j.Offsets.Min,
+			Max: j.Offsets.Max,
+		},
+	}
+}

pkg/blockbuilder/types/interfaces.go

@@ -24,6 +24,15 @@ type Scheduler interface {
 
 // Transport defines the interface for communication between block builders and scheduler
 type Transport interface {
+	BuilderTransport
+	SchedulerTransport
+}
+
+// SchedulerTransport is for calls originating from the scheduler
+type SchedulerTransport interface{}
+
+// BuilderTransport is for calls originating from the builder
+type BuilderTransport interface {
 	// SendGetJobRequest sends a request to get a new job
 	SendGetJobRequest(ctx context.Context, req *GetJobRequest) (*GetJobResponse, error)
 	// SendCompleteJob sends a job completion notification

pkg/blockbuilder/types/job.go

@@ -4,8 +4,7 @@ import "fmt"
 
 // Job represents a block building task.
 type Job struct {
-	ID     string
-	Status JobStatus
+	ID string
 	// Partition and offset information
 	Partition int
 	Offsets   Offsets
@@ -30,7 +29,6 @@ type Offsets struct {
 func NewJob(partition int, offsets Offsets) *Job {
 	return &Job{
 		ID:        GenerateJobID(partition, offsets),
-		Status:    JobStatusPending,
 		Partition: partition,
 		Offsets:   offsets,
 	}

pkg/blockbuilder/types/proto/blockbuilder.proto

@@ -0,0 +1,57 @@
+syntax = "proto3";
+
+package blockbuilder.types;
+
+option go_package = "github.com/grafana/loki/pkg/blockbuilder/types/proto";
+
+// BlockBuilderService defines the gRPC service for block builder communication
+service BlockBuilderService {
+  // GetJob requests a new job from the scheduler
+  rpc GetJob(GetJobRequest) returns (GetJobResponse) {}
+  // CompleteJob notifies the scheduler that a job has been completed
+  rpc CompleteJob(CompleteJobRequest) returns (CompleteJobResponse) {}
+  // SyncJob syncs job state with the scheduler
+  rpc SyncJob(SyncJobRequest) returns (SyncJobResponse) {}
+}
+
+// GetJobRequest represents a request for a new job
+message GetJobRequest {
+  string builder_id = 1;
+}
+
+// GetJobResponse contains the response for a job request
+message GetJobResponse {
+  Job job = 1;
+  bool ok = 2;
+}
+
+// CompleteJobRequest represents a job completion notification
+message CompleteJobRequest {
+  string builder_id = 1;
+  Job job = 2;
+}
+
+// CompleteJobResponse is an empty response for job completion
+message CompleteJobResponse {}
+
+// SyncJobRequest represents a job sync request
+message SyncJobRequest {
+  string builder_id = 1;
+  Job job = 2;
+}
+
+// SyncJobResponse is an empty response for job sync
+message SyncJobResponse {}
+
+// Offsets represents the start and end offsets for a job
+message Offsets {
+  int64 min = 1;
+  int64 max = 2;
+}
+
+// Job represents a block building job
+message Job {
+  string id = 1;
+  int32 partition = 2;
+  Offsets offsets = 3;
+}

pkg/blockbuilder/types/transport.go

@@ -1,58 +1,56 @@
-package builder
+package types
 
 import (
 	"context"
-
-	"github.com/grafana/loki/v3/pkg/blockbuilder/types"
 )
 
 var (
-	_ types.Transport = unimplementedTransport{}
-	_ types.Transport = &MemoryTransport{}
+	_ Transport = unimplementedTransport{}
+	_ Transport = &MemoryTransport{}
 )
 
 // unimplementedTransport provides default implementations that panic
 type unimplementedTransport struct{}
 
-func (t unimplementedTransport) SendGetJobRequest(_ context.Context, _ *types.GetJobRequest) (*types.GetJobResponse, error) {
+func (t unimplementedTransport) SendGetJobRequest(_ context.Context, _ *GetJobRequest) (*GetJobResponse, error) {
 	panic("unimplemented")
 }
 
-func (t unimplementedTransport) SendCompleteJob(_ context.Context, _ *types.CompleteJobRequest) error {
+func (t unimplementedTransport) SendCompleteJob(_ context.Context, _ *CompleteJobRequest) error {
 	panic("unimplemented")
 }
 
-func (t unimplementedTransport) SendSyncJob(_ context.Context, _ *types.SyncJobRequest) error {
+func (t unimplementedTransport) SendSyncJob(_ context.Context, _ *SyncJobRequest) error {
 	panic("unimplemented")
 }
 
 // MemoryTransport implements Transport interface for in-memory communication
 type MemoryTransport struct {
-	scheduler types.Scheduler
+	scheduler Scheduler
 }
 
 // NewMemoryTransport creates a new in-memory transport instance
-func NewMemoryTransport(scheduler types.Scheduler) *MemoryTransport {
+func NewMemoryTransport(scheduler Scheduler) *MemoryTransport {
 	return &MemoryTransport{
 		scheduler: scheduler,
 	}
 }
 
-func (t *MemoryTransport) SendGetJobRequest(ctx context.Context, req *types.GetJobRequest) (*types.GetJobResponse, error) {
+func (t *MemoryTransport) SendGetJobRequest(ctx context.Context, req *GetJobRequest) (*GetJobResponse, error) {
 	job, ok, err := t.scheduler.HandleGetJob(ctx, req.BuilderID)
 	if err != nil {
 		return nil, err
 	}
-	return &types.GetJobResponse{
+	return &GetJobResponse{
 		Job: job,
 		OK:  ok,
 	}, nil
 }
 
-func (t *MemoryTransport) SendCompleteJob(ctx context.Context, req *types.CompleteJobRequest) error {
+func (t *MemoryTransport) SendCompleteJob(ctx context.Context, req *CompleteJobRequest) error {
 	return t.scheduler.HandleCompleteJob(ctx, req.BuilderID, req.Job)
 }
 
-func (t *MemoryTransport) SendSyncJob(ctx context.Context, req *types.SyncJobRequest) error {
+func (t *MemoryTransport) SendSyncJob(ctx context.Context, req *SyncJobRequest) error {
 	return t.scheduler.HandleSyncJob(ctx, req.BuilderID, req.Job)
 }