more otel

pkg/actor/mutation_registry.go

@@ -1,12 +1,14 @@
 package actor
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"reflect"
 	"sync"
 
 	"github.com/gogo/protobuf/proto"
+	"go.opentelemetry.io/otel/attribute"
 )
 
 type ApplyResult struct {
@@ -16,27 +18,25 @@ type ApplyResult struct {
 }
 
 type MutationProcessor interface {
-	Process(grain any) error
+	Process(ctx context.Context, grain any) error
 }
 
 type BasicMutationProcessor[V any] struct {
-	processor func(any) error
+	processor func(ctx context.Context, grain V) error
 }
 
-func NewMutationProcessor[V any](process func(V) error) MutationProcessor {
+func NewMutationProcessor[V any](process func(ctx context.Context, grain V) error) MutationProcessor {
 	return &BasicMutationProcessor[V]{
-		processor: func(v any) error {
-			return process(v.(V))
-		},
+		processor: process,
 	}
 }
 
-func (p *BasicMutationProcessor[V]) Process(grain any) error {
-	return p.processor(grain)
+func (p *BasicMutationProcessor[V]) Process(ctx context.Context, grain any) error {
+	return p.processor(ctx, grain.(V))
 }
 
 type MutationRegistry interface {
-	Apply(grain any, msg ...proto.Message) ([]ApplyResult, error)
+	Apply(ctx context.Context, grain any, msg ...proto.Message) ([]ApplyResult, error)
 	RegisterMutations(handlers ...MutationHandler)
 	Create(typeName string) (proto.Message, bool)
 	GetTypeName(msg proto.Message) (string, bool)
@@ -192,7 +192,15 @@ func (r *ProtoMutationRegistry) Create(typeName string) (proto.Message, bool) {
 // Returns updated grain if successful.
 //
 // If the mutation is not registered, returns (nil, ErrMutationNotRegistered).
-func (r *ProtoMutationRegistry) Apply(grain any, msg ...proto.Message) ([]ApplyResult, error) {
+func (r *ProtoMutationRegistry) Apply(ctx context.Context, grain any, msg ...proto.Message) ([]ApplyResult, error) {
+
+	parentCtx, span := tracer.Start(ctx, "apply mutations")
+	defer span.End()
+	span.SetAttributes(
+		attribute.String("component", "registry"),
+		attribute.Int("mutations", len(msg)),
+	)
+
 	results := make([]ApplyResult, 0, len(msg))
 
 	if grain == nil {
@@ -214,20 +222,30 @@ func (r *ProtoMutationRegistry) Apply(grain any, msg ...proto.Message) ([]ApplyR
 			continue
 		}
 		rt := indirectType(reflect.TypeOf(m))
+		_, msgSpan := tracer.Start(parentCtx, rt.Name())
+
 		r.mutationRegistryMu.RLock()
 		entry, ok := r.mutationRegistry[rt]
 		r.mutationRegistryMu.RUnlock()
 		if !ok {
 			results = append(results, ApplyResult{Error: ErrMutationNotRegistered, Type: rt.Name(), Mutation: m})
 			continue
+		} else {
+			err := entry.Handle(grain, m)
+			if err != nil {
+				msgSpan.RecordError(err)
+			}
+			results = append(results, ApplyResult{Error: err, Type: rt.Name(), Mutation: m})
 		}
-		err := entry.Handle(grain, m)
-		results = append(results, ApplyResult{Error: err, Type: rt.Name(), Mutation: m})
+		msgSpan.End()
 	}
 
 	if len(results) > 0 {
+		processCtx, processSpan := tracer.Start(ctx, "after mutation processors")
+		defer processSpan.End()
 		for _, processor := range r.processors {
-			err := processor.Process(grain)
+
+			err := processor.Process(processCtx, grain)
 			if err != nil {
 				return results, err
 			}