Eng 565 messagehandler metrics (#4833)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
	- Enhanced telemetry metrics for message handling
	- Added new performance tracking metrics for message processing

- **Refactor**
	- Updated method signatures to support metric recording
- Restructured metrics tracking from worker-related to message-handling
focused

- **Chores**
	- Improved observability of message processing operations
	- Reorganized constant attributes for better metric management

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: wdbaruni

pkg/orchestrator/message_handler.go

@@ -8,11 +8,13 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/rs/zerolog/log"
+	"go.opentelemetry.io/otel/attribute"
 
 	"github.com/bacalhau-project/bacalhau/pkg/jobstore"
 	"github.com/bacalhau-project/bacalhau/pkg/lib/envelope"
 	"github.com/bacalhau-project/bacalhau/pkg/models"
 	"github.com/bacalhau-project/bacalhau/pkg/models/messages"
+	"github.com/bacalhau-project/bacalhau/pkg/telemetry"
 )
 
 // MessageHandler base implementation of requester Endpoint
@@ -35,27 +37,38 @@ func (m *MessageHandler) ShouldProcess(ctx context.Context, message *envelope.Me
 
 // HandleMessage handles incoming messages
 // TODO: handle messages arriving out of order gracefully
-func (m *MessageHandler) HandleMessage(ctx context.Context, message *envelope.Message) error {
-	var err error
+func (m *MessageHandler) HandleMessage(ctx context.Context, message *envelope.Message) (err error) {
+	metrics := telemetry.NewMetricRecorder(
+		attribute.String(AttrMessageType, message.Metadata.Get(envelope.KeyMessageType)),
+		attribute.String(AttrOutcomeKey, AttrOutcomeSuccess),
+	)
+	defer func() {
+		metrics.Count(ctx, messageHandlerProcessCount)
+		metrics.Done(ctx, messageHandlerProcessDuration)
+	}()
+
 	switch message.Metadata.Get(envelope.KeyMessageType) {
 	case messages.BidResultMessageType:
-		err = m.OnBidComplete(ctx, message)
+		err = m.OnBidComplete(ctx, metrics, message)
 	case messages.RunResultMessageType:
-		err = m.OnRunComplete(ctx, message)
+		err = m.OnRunComplete(ctx, metrics, message)
 	case messages.ComputeErrorMessageType:
-		err = m.OnComputeFailure(ctx, message)
+		err = m.OnComputeFailure(ctx, metrics, message)
 	}
 
-	return m.handleError(ctx, message, err)
+	return m.handleError(ctx, metrics, message, err)
 }
 
 // handleError logs the error with context and returns nil.
 // In the future, this can be extended to handle different error types differently.
-func (m *MessageHandler) handleError(ctx context.Context, message *envelope.Message, err error) error {
+func (m *MessageHandler) handleError(ctx context.Context, metrics *telemetry.MetricRecorder, message *envelope.Message, err error) error {
 	if err == nil {
 		return nil
 	}
 
+	metrics.Error(err)
+	metrics.AddAttributes(attribute.String(AttrOutcomeKey, AttrOutcomeFailure))
+
 	// For now, just log the error and return nil
 	logger := log.Ctx(ctx).Error()
 	for key, value := range message.Metadata.ToMap() {
@@ -66,7 +79,7 @@ func (m *MessageHandler) handleError(ctx context.Context, message *envelope.Mess
 }
 
 // OnBidComplete handles the completion of a bid request
-func (m *MessageHandler) OnBidComplete(ctx context.Context, message *envelope.Message) error {
+func (m *MessageHandler) OnBidComplete(ctx context.Context, metrics *telemetry.MetricRecorder, message *envelope.Message) error {
 	result, ok := message.Payload.(*messages.BidResult)
 	if !ok {
 		return envelope.NewErrUnexpectedPayloadType("BidResult", reflect.TypeOf(message.Payload).String())
@@ -92,6 +105,7 @@ func (m *MessageHandler) OnBidComplete(ctx context.Context, message *envelope.Me
 	}
 
 	txContext, err := m.store.BeginTx(ctx)
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartBeginTx)
 	if err != nil {
 		return fmt.Errorf("failed to begin transaction: %w", err)
 	}
@@ -101,30 +115,36 @@ func (m *MessageHandler) OnBidComplete(ctx context.Context, message *envelope.Me
 	if err = m.store.UpdateExecution(txContext, updateRequest); err != nil {
 		return err
 	}
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartUpdateExec)
 
 	// enqueue evaluation to allow the scheduler to either accept the bid, or find a new node
 	err = m.enqueueEvaluation(txContext, result.JobID, result.JobType)
 	if err != nil {
 		return err
 	}
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartCreateEval)
 
-	return txContext.Commit()
+	err = txContext.Commit()
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartCommitTx)
+	return err
 }
 
-func (m *MessageHandler) OnRunComplete(ctx context.Context, message *envelope.Message) error {
+func (m *MessageHandler) OnRunComplete(ctx context.Context, metrics *telemetry.MetricRecorder, message *envelope.Message) error {
 	result, ok := message.Payload.(*messages.RunResult)
 	if !ok {
 		return envelope.NewErrUnexpectedPayloadType("RunResult", reflect.TypeOf(message.Payload).String())
 	}
 
 	txContext, err := m.store.BeginTx(ctx)
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartBeginTx)
 	if err != nil {
 		return fmt.Errorf("failed to begin transaction: %w", err)
 	}
 
 	defer txContext.Rollback() //nolint:errcheck
 
 	job, err := m.store.GetJob(txContext, result.JobID)
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartGetJob)
 	if err != nil {
 		return err
 	}
@@ -158,22 +178,27 @@ func (m *MessageHandler) OnRunComplete(ctx context.Context, message *envelope.Me
 	if err = m.store.UpdateExecution(txContext, updateRequest); err != nil {
 		return err
 	}
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartUpdateExec)
 
 	// enqueue evaluation to allow the scheduler to mark the job as completed if all executions are completed
 	if err = m.enqueueEvaluation(txContext, result.JobID, result.JobType); err != nil {
 		return err
 	}
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartCreateEval)
 
-	return txContext.Commit()
+	err = txContext.Commit()
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartCommitTx)
+	return err
 }
 
-func (m *MessageHandler) OnComputeFailure(ctx context.Context, message *envelope.Message) error {
+func (m *MessageHandler) OnComputeFailure(ctx context.Context, metrics *telemetry.MetricRecorder, message *envelope.Message) error {
 	result, ok := message.Payload.(*messages.ComputeError)
 	if !ok {
 		return envelope.NewErrUnexpectedPayloadType("ComputeError", reflect.TypeOf(message.Payload).String())
 	}
 
 	txContext, err := m.store.BeginTx(ctx)
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartBeginTx)
 	if err != nil {
 		return fmt.Errorf("failed to begin transaction: %w", err)
 	}
@@ -197,13 +222,17 @@ func (m *MessageHandler) OnComputeFailure(ctx context.Context, message *envelope
 	}); err != nil {
 		return err
 	}
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartUpdateExec)
 
 	// enqueue evaluation to allow the scheduler find other nodes, or mark the job as failed
 	if err = m.enqueueEvaluation(txContext, result.JobID, result.JobType); err != nil {
 		return err
 	}
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartCreateEval)
 
-	return txContext.Commit()
+	err = txContext.Commit()
+	metrics.Latency(ctx, messageHandlerProcessPartDuration, AttrPartCommitTx)
+	return err
 }
 
 // enqueueEvaluation enqueues an evaluation to allow the scheduler to either accept the bid, or find a new node

pkg/orchestrator/message_handler_test.go

@@ -42,7 +42,7 @@ func (suite *MessageHandlerTestSuite) TestShouldProcess() {
 	suite.False(suite.handler.ShouldProcess(context.Background(), envelope.NewMessage(nil).WithMetadataValue(envelope.KeyMessageType, "UnknownType")))
 }
 
-func (suite *MessageHandlerTestSuite) TestOnBidComplete_Accepted() {
+func (suite *MessageHandlerTestSuite) TestHandleBidAccepted() {
 	ctx := context.Background()
 	bidResult := &messages.BidResult{
 		BaseResponse: messages.BaseResponse{
@@ -60,11 +60,11 @@ func (suite *MessageHandlerTestSuite) TestOnBidComplete_Accepted() {
 	suite.mockTx.EXPECT().Commit().Return(nil)
 	suite.mockTx.EXPECT().Rollback().Return(nil)
 
-	err := suite.handler.OnBidComplete(ctx, message)
+	err := suite.handler.HandleMessage(ctx, message)
 	suite.NoError(err)
 }
 
-func (suite *MessageHandlerTestSuite) TestOnBidComplete_Rejected() {
+func (suite *MessageHandlerTestSuite) TestHandleBidRejected() {
 	ctx := context.Background()
 	bidResult := &messages.BidResult{
 		BaseResponse: messages.BaseResponse{
@@ -82,11 +82,11 @@ func (suite *MessageHandlerTestSuite) TestOnBidComplete_Rejected() {
 	suite.mockTx.EXPECT().Commit().Return(nil)
 	suite.mockTx.EXPECT().Rollback().Return(nil)
 
-	err := suite.handler.OnBidComplete(ctx, message)
+	err := suite.handler.HandleMessage(ctx, message)
 	suite.NoError(err)
 }
 
-func (suite *MessageHandlerTestSuite) TestOnRunComplete() {
+func (suite *MessageHandlerTestSuite) TestHandleRunComplete() {
 	ctx := context.Background()
 	runResult := &messages.RunResult{
 		BaseResponse: messages.BaseResponse{
@@ -106,11 +106,35 @@ func (suite *MessageHandlerTestSuite) TestOnRunComplete() {
 	suite.mockTx.EXPECT().Commit().Return(nil)
 	suite.mockTx.EXPECT().Rollback().Return(nil)
 
-	err := suite.handler.OnRunComplete(ctx, message)
+	err := suite.handler.HandleMessage(ctx, message)
 	suite.NoError(err)
 }
 
-func (suite *MessageHandlerTestSuite) TestOnComputeFailure() {
+func (suite *MessageHandlerTestSuite) TestHandleRunCompleteForLongRunningJob() {
+	ctx := context.Background()
+	runResult := &messages.RunResult{
+		BaseResponse: messages.BaseResponse{
+			ExecutionID: "exec-1",
+			JobID:       "job-1",
+			JobType:     "service",
+		},
+		PublishResult:    &models.SpecConfig{Type: "ipfs"},
+		RunCommandResult: &models.RunCommandResult{ExitCode: 0},
+	}
+	message := envelope.NewMessage(runResult).WithMetadataValue(envelope.KeyMessageType, messages.RunResultMessageType)
+
+	suite.mockStore.EXPECT().BeginTx(gomock.Any()).Return(suite.mockTx, nil)
+	suite.mockStore.EXPECT().GetJob(suite.mockTx, "job-1").Return(models.Job{Type: "service"}, nil)
+	suite.mockStore.EXPECT().UpdateExecution(suite.mockTx, gomock.Any()).Return(nil)
+	suite.mockStore.EXPECT().CreateEvaluation(suite.mockTx, gomock.Any()).Return(nil)
+	suite.mockTx.EXPECT().Commit().Return(nil)
+	suite.mockTx.EXPECT().Rollback().Return(nil)
+
+	err := suite.handler.HandleMessage(ctx, message)
+	suite.NoError(err)
+}
+
+func (suite *MessageHandlerTestSuite) TestHandleComputeFailure() {
 	ctx := context.Background()
 	computeError := &messages.ComputeError{
 		BaseResponse: messages.BaseResponse{
@@ -127,11 +151,11 @@ func (suite *MessageHandlerTestSuite) TestOnComputeFailure() {
 	suite.mockTx.EXPECT().Commit().Return(nil)
 	suite.mockTx.EXPECT().Rollback().Return(nil)
 
-	err := suite.handler.OnComputeFailure(ctx, message)
+	err := suite.handler.HandleMessage(ctx, message)
 	suite.NoError(err)
 }
 
-func (suite *MessageHandlerTestSuite) TestOnBidComplete_PropagatesErrors() {
+func (suite *MessageHandlerTestSuite) TestHandleMessagePropagatesErrors() {
 	ctx := context.Background()
 	bidResult := &messages.BidResult{
 		BaseResponse: messages.BaseResponse{
@@ -148,8 +172,8 @@ func (suite *MessageHandlerTestSuite) TestOnBidComplete_PropagatesErrors() {
 	suite.mockStore.EXPECT().UpdateExecution(suite.mockTx, gomock.Any()).Return(expectedErr)
 	suite.mockTx.EXPECT().Rollback().Return(nil)
 
-	err := suite.handler.OnBidComplete(ctx, message)
-	suite.ErrorIs(err, expectedErr)
+	err := suite.handler.HandleMessage(ctx, message)
+	suite.NoError(err) // HandleMessage swallows errors after logging them
 }
 
 func TestMessageHandlerTestSuite(t *testing.T) {

pkg/orchestrator/metrics.go

@@ -1,37 +1,15 @@
 package orchestrator
 
 import (
-	"github.com/bacalhau-project/bacalhau/pkg/telemetry"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
-)
 
-var (
-	Meter = otel.GetMeterProvider().Meter("orchestrator")
+	"github.com/bacalhau-project/bacalhau/pkg/telemetry"
 )
 
-// Metrics for monitoring worker
 var (
-	WorkerDequeueFaults = telemetry.Must(Meter.Int64Counter(
-		"worker_dequeue_faults",
-		metric.WithDescription("Number of times a worker failed to dequeue an evaluation"),
-	))
-
-	WorkerProcessFaults = telemetry.Must(Meter.Int64Counter(
-		"worker_process_faults",
-		metric.WithDescription("Number of times a worker failed to process an evaluation"),
-	))
-
-	WorkerAckFaults = telemetry.Must(Meter.Int64Counter(
-		"worker_ack_faults",
-		metric.WithDescription("Number of times a worker failed to ack an evaluation back to the broker"),
-	))
-
-	WorkerNackFaults = telemetry.Must(Meter.Int64Counter(
-		"worker_nack_faults",
-		metric.WithDescription("Number of times a worker failed to nack an evaluation back to the broker"),
-	))
+	Meter = otel.GetMeterProvider().Meter("orchestrator")
 )
 
 // Metrics for monitoring evaluation broker
@@ -62,6 +40,45 @@ var (
 	))
 )
 
+// Message handler metrics
+var (
+	// Message processing metrics
+	messageHandlerProcessDuration = telemetry.Must(Meter.Float64Histogram(
+		"message.handler.process.duration",
+		metric.WithDescription("Time taken to process a single message"),
+		metric.WithUnit("s"),
+		metric.WithExplicitBucketBoundaries(telemetry.DurationMsBuckets...),
+	))
+
+	messageHandlerProcessPartDuration = telemetry.Must(Meter.Float64Histogram(
+		"message.handler.process.part.duration",
+		metric.WithDescription("Time taken for sub-operations within message handling"),
+		metric.WithUnit("s"),
+		metric.WithExplicitBucketBoundaries(telemetry.DurationMsBuckets...),
+	))
+
+	messageHandlerProcessCount = telemetry.Must(Meter.Int64Counter(
+		"message.handler.process.count",
+		metric.WithDescription("Number of messages processed"),
+		metric.WithUnit("1"),
+	))
+)
+
+const (
+	AttrEvalType    = "eval_type"
+	AttrMessageType = "message_type"
+
+	AttrPartBeginTx    = "begin_transaction"
+	AttrPartGetJob     = "get_job"
+	AttrPartCommitTx   = "commit_transaction"
+	AttrPartUpdateExec = "update_execution"
+	AttrPartCreateEval = "create_evaluation"
+
+	AttrOutcomeKey     = "outcome"
+	AttrOutcomeSuccess = "success"
+	AttrOutcomeFailure = "failure"
+)
+
 func EvalTypeAttribute(evaluationType string) metric.MeasurementOption {
-	return metric.WithAttributes(attribute.String("eval_type", evaluationType))
+	return metric.WithAttributes(attribute.String(AttrEvalType, evaluationType))
 }