chore: Remove RunSpan in favor of defer style spans

pkg/authorization/gateway_authorizer.go

@@ -6,7 +6,6 @@ import (
 
 	log "github.com/sirupsen/logrus"
 	"github.com/superplanehq/superplane/pkg/models"
-	"github.com/superplanehq/superplane/pkg/telemetry"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
@@ -51,12 +50,7 @@ func (a *GatewayAuthorizer) AuthorizeHTTP(
 		return nil, status.Error(codes.NotFound, "Not found")
 	}
 
-	var allowed bool
-	err := telemetry.RunSpan(ctx, "auth.check_permission", func(ctx context.Context) error {
-		var checkErr error
-		allowed, checkErr = checkOrganizationPermission(ctx, a.auth, userID, organizationID, rule.Resource, rule.Action)
-		return checkErr
-	})
+	allowed, err := checkOrganizationPermission(ctx, a.auth, userID, organizationID, rule.Resource, rule.Action)
 	if err != nil {
 		return nil, err
 	}

pkg/authorization/organization_permissions.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"sync"
 	"time"
+
+	"github.com/superplanehq/superplane/pkg/telemetry"
 )
 
 const organizationPermissionCacheTTL = 15 * time.Minute
@@ -62,6 +64,10 @@ type organizationPermissionChecker interface {
 }
 
 func checkOrganizationPermission(ctx context.Context, auth organizationPermissionChecker, userID, orgID, resource, action string) (bool, error) {
+	var err error
+	ctx, done := telemetry.Span(ctx, "auth.check_organization_permission")
+	defer done(&err)
+
 	if allowed, ok := cachedOrganizationPermissionAllowed(userID, orgID, resource, action); ok && allowed {
 		return true, nil
 	}

pkg/grpc/actions/canvases/common.go

@@ -6,7 +6,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/superplanehq/superplane/pkg/grpc/actions/canvases/changesets"
-	"github.com/superplanehq/superplane/pkg/grpc/errors"
+	grpcerrors "github.com/superplanehq/superplane/pkg/grpc/errors"
 	"github.com/superplanehq/superplane/pkg/models"
 	pb "github.com/superplanehq/superplane/pkg/protos/canvases"
 	"github.com/superplanehq/superplane/pkg/telemetry"
@@ -15,87 +15,68 @@ import (
 
 const canvasNameAlreadyExistsMessage = "Canvas with the same name already exists"
 
-func checkCanvasExistence(ctx context.Context, db *gorm.DB, orgID, canvasID uuid.UUID) error {
-	return telemetry.RunSpan(ctx, "canvases.check_canvas_existence", func(ctx context.Context) error {
-		exists, err := models.CheckCanvasExistence(db, orgID, canvasID)
-		if err != nil {
-			return err
-		}
-		if !exists {
-			return gorm.ErrRecordNotFound
-		}
+func checkCanvasExistence(ctx context.Context, db *gorm.DB, orgID, canvasID uuid.UUID) (err error) {
+	ctx, done := telemetry.Span(ctx, "canvases.check_canvas_existence")
+	defer done(&err)
 
-		return nil
-	})
+	exists, err := models.CheckCanvasExistence(db.WithContext(ctx), orgID, canvasID)
+	if err != nil {
+		return err
+	}
+	if !exists {
+		return gorm.ErrRecordNotFound
+	}
+
+	return nil
+}
+
+func loadCanvas(ctx context.Context, db *gorm.DB, orgID, canvasID uuid.UUID) (canvas *models.Canvas, err error) {
+	ctx, done := telemetry.Span(ctx, "canvases.find_canvas")
+	defer done(&err)
+
+	return models.FindCanvasInTransaction(db.WithContext(ctx), orgID, canvasID)
 }
 
-func loadCanvas(ctx context.Context, db *gorm.DB, orgID, canvasID uuid.UUID) (*models.Canvas, error) {
-	var canvas *models.Canvas
-	err := telemetry.RunSpan(ctx, "canvases.find_canvas", func(ctx context.Context) error {
-		var findErr error
-		canvas, findErr = models.FindCanvasInTransaction(db, orgID, canvasID)
-		return findErr
-	})
+func loadLiveCanvasVersion(ctx context.Context, db *gorm.DB, canvas *models.Canvas) (liveVersion *models.CanvasVersion, err error) {
+	ctx, done := telemetry.Span(ctx, "canvases.load_live_version")
+	defer done(&err)
+
+	return models.FindLiveCanvasVersionByCanvasInTransaction(db.WithContext(ctx), canvas)
+}
+
+func loadCanvasStatus(ctx context.Context, db *gorm.DB, canvasID uuid.UUID) (canvasStatus *pb.Canvas_Status, err error) {
+	ctx, done := telemetry.Span(ctx, "canvases.load_status")
+	defer done(&err)
+
+	lastExecutions, err := models.FindLastExecutionPerNode(db.WithContext(ctx), canvasID)
 	if err != nil {
 		return nil, err
 	}
 
-	return canvas, nil
-}
+	executionResources, err := LoadNodeExecutionResources(db.WithContext(ctx), lastExecutions)
+	if err != nil {
+		return nil, err
+	}
 
-func loadLiveCanvasVersion(ctx context.Context, db *gorm.DB, canvas *models.Canvas) (*models.CanvasVersion, error) {
-	var liveVersion *models.CanvasVersion
-	err := telemetry.RunSpan(ctx, "canvases.load_live_version", func(ctx context.Context) error {
-		var loadErr error
-		liveVersion, loadErr = models.FindLiveCanvasVersionByCanvasInTransaction(db, canvas)
-		return loadErr
-	})
+	serializedExecutions, err := SerializeNodeExecutions(lastExecutions, executionResources)
 	if err != nil {
 		return nil, err
 	}
 
-	return liveVersion, nil
-}
+	lastEvents, err := models.FindLastEventPerNode(db.WithContext(ctx), canvasID)
+	if err != nil {
+		return nil, err
+	}
 
-func loadCanvasStatus(ctx context.Context, db *gorm.DB, canvasID uuid.UUID) (*pb.Canvas_Status, error) {
-	var canvasStatus *pb.Canvas_Status
-	err := telemetry.RunSpan(ctx, "canvases.load_status", func(ctx context.Context) error {
-		lastExecutions, loadErr := models.FindLastExecutionPerNode(db, canvasID)
-		if loadErr != nil {
-			return loadErr
-		}
-
-		executionResources, loadErr := LoadNodeExecutionResources(db, lastExecutions)
-		if loadErr != nil {
-			return loadErr
-		}
-
-		serializedExecutions, loadErr := SerializeNodeExecutions(lastExecutions, executionResources)
-		if loadErr != nil {
-			return loadErr
-		}
-
-		lastEvents, loadErr := models.FindLastEventPerNode(db, canvasID)
-		if loadErr != nil {
-			return loadErr
-		}
-
-		serializedEvents, loadErr := SerializeCanvasEvents(lastEvents)
-		if loadErr != nil {
-			return loadErr
-		}
-
-		canvasStatus = &pb.Canvas_Status{
-			LastExecutions: serializedExecutions,
-			LastEvents:     serializedEvents,
-		}
-		return nil
-	})
+	serializedEvents, err := SerializeCanvasEvents(lastEvents)
 	if err != nil {
 		return nil, err
 	}
 
-	return canvasStatus, nil
+	return &pb.Canvas_Status{
+		LastExecutions: serializedExecutions,
+		LastEvents:     serializedEvents,
+	}, nil
 }
 
 func ensureCanvasNameAvailableInTransaction(

pkg/grpc/actions/canvases/describe_canvas.go

@@ -11,8 +11,16 @@ import (
 	pb "github.com/superplanehq/superplane/pkg/protos/canvases"
 	"github.com/superplanehq/superplane/pkg/registry"
 	"github.com/superplanehq/superplane/pkg/telemetry"
+	"gorm.io/gorm"
 )
 
+func loadCanvasCreator(ctx context.Context, db *gorm.DB, canvas *models.Canvas) (user *models.User, err error) {
+	ctx, done := telemetry.Span(ctx, "canvases.load_creator")
+	defer done(&err)
+
+	return models.FindMaybeDeletedUserByIDInTransaction(db.WithContext(ctx), canvas.OrganizationID.String(), canvas.CreatedBy.String())
+}
+
 func DescribeCanvas(ctx context.Context, registry *registry.Registry, organizationID string, id string) (*pb.DescribeCanvasResponse, error) {
 	canvasID, err := uuid.Parse(id)
 	if err != nil {
@@ -29,11 +37,7 @@ func DescribeCanvas(ctx context.Context, registry *registry.Registry, organizati
 
 	var user *models.User
 	if canvas.CreatedBy != nil {
-		err = telemetry.RunSpan(ctx, "canvases.load_creator", func(ctx context.Context) error {
-			var loadErr error
-			user, loadErr = models.FindMaybeDeletedUserByIDInTransaction(db, canvas.OrganizationID.String(), canvas.CreatedBy.String())
-			return loadErr
-		})
+		user, err = loadCanvasCreator(ctx, db, canvas)
 		if err != nil {
 			return nil, err
 		}

pkg/grpc/actions/canvases/list_canvas_memories.go

@@ -36,12 +36,7 @@ func ListCanvasMemories(ctx context.Context, registry *registry.Registry, organi
 		return nil, grpcerrors.Internal(err, "failed to load canvas")
 	}
 
-	var records []models.CanvasMemory
-	err = telemetry.RunSpan(ctx, "memories.list", func(ctx context.Context) error {
-		var listErr error
-		records, listErr = models.ListCanvasMemoriesInTransaction(database.DB(ctx), canvasUUID)
-		return listErr
-	})
+	records, err := listCanvasMemories(ctx, canvasUUID)
 	if err != nil {
 		return nil, grpcerrors.Internal(err, "failed to list canvas memories")
 	}
@@ -56,27 +51,29 @@ func ListCanvasMemories(ctx context.Context, registry *registry.Registry, organi
 	}, nil
 }
 
-func serializeCanvasMemories(ctx context.Context, records []models.CanvasMemory) ([]*pb.CanvasMemory, error) {
-	var items []*pb.CanvasMemory
-	err := telemetry.RunSpan(ctx, "memories.serialize", func(ctx context.Context) error {
-		items = make([]*pb.CanvasMemory, 0, len(records))
-		for _, record := range records {
-			item, itemErr := canvasMemoryToProto(record)
-			if itemErr != nil {
-				return itemErr
-			}
-
-			items = append(items, item)
-		}
+func listCanvasMemories(ctx context.Context, canvasUUID uuid.UUID) (records []models.CanvasMemory, err error) {
+	ctx, done := telemetry.Span(ctx, "memories.list")
+	defer done(&err)
+
+	return models.ListCanvasMemoriesInTransaction(database.DB(ctx), canvasUUID)
+}
+
+func serializeCanvasMemories(ctx context.Context, records []models.CanvasMemory) (items []*pb.CanvasMemory, err error) {
+	ctx, done := telemetry.Span(ctx, "memories.serialize")
+	defer done(&err)
 
-		if span := trace.SpanFromContext(ctx); span.IsRecording() {
-			span.SetAttributes(attribute.Int("memories.count", len(records)))
+	items = make([]*pb.CanvasMemory, 0, len(records))
+	for _, record := range records {
+		item, itemErr := canvasMemoryToProto(record)
+		if itemErr != nil {
+			return nil, itemErr
 		}
 
-		return nil
-	})
-	if err != nil {
-		return nil, err
+		items = append(items, item)
+	}
+
+	if span := trace.SpanFromContext(ctx); span.IsRecording() {
+		span.SetAttributes(attribute.Int("memories.count", len(records)))
 	}
 
 	return items, nil

pkg/grpc/actions/canvases/list_canvas_versions.go

@@ -2,6 +2,7 @@ package canvases
 
 import (
 	"context"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/superplanehq/superplane/pkg/authentication"
@@ -54,25 +55,7 @@ func ListCanvasVersionsPaginated(
 	limit = getCanvasVersionLimit(limit)
 	beforeTime := getBefore(before)
 
-	var publishedVersions []models.CanvasVersion
-	var publishedCount int64
-	err = telemetry.RunSpan(ctx, "canvases.list_published_versions", func(ctx context.Context) error {
-		return database.DB(ctx).Transaction(func(tx *gorm.DB) error {
-			versions, versionsErr := models.ListPublishedCanvasVersionsInTransaction(tx, canvasUUID, int(limit), beforeTime)
-			if versionsErr != nil {
-				return versionsErr
-			}
-			publishedVersions = versions
-
-			count, countErr := models.CountPublishedCanvasVersionsInTransaction(tx, canvasUUID)
-			if countErr != nil {
-				return countErr
-			}
-			publishedCount = count
-
-			return nil
-		})
-	})
+	publishedVersions, publishedCount, err := listPublishedCanvasVersions(ctx, canvasUUID, int(limit), beforeTime)
 	if err != nil {
 		return nil, grpcerrors.Internal(err, "failed to list canvas versions")
 	}
@@ -98,25 +81,7 @@ func listDraftCanvasVersions(
 	limit = getCanvasVersionLimit(limit)
 	beforeTime := getBefore(before)
 
-	var draftVersions []models.CanvasVersion
-	var draftCount int64
-	err := telemetry.RunSpan(ctx, "canvases.list_draft_versions", func(ctx context.Context) error {
-		return database.DB(ctx).Transaction(func(tx *gorm.DB) error {
-			versions, versionsErr := models.ListDraftBranchesForCanvasInTransaction(tx, canvasID, ownerID, int(limit), beforeTime)
-			if versionsErr != nil {
-				return versionsErr
-			}
-			draftVersions = versions
-
-			count, countErr := models.CountDraftBranchesForCanvasInTransaction(tx, canvasID, ownerID)
-			if countErr != nil {
-				return countErr
-			}
-			draftCount = count
-
-			return nil
-		})
-	})
+	draftVersions, draftCount, err := listDraftCanvasVersionsInTransaction(ctx, canvasID, ownerID, int(limit), beforeTime)
 	if err != nil {
 		return nil, grpcerrors.Internal(err, "failed to list canvas versions")
 	}
@@ -168,3 +133,39 @@ func getLastDraftCanvasVersionTimestamp(versions []models.CanvasVersion) *timest
 
 	return timestamppb.New(*lastVersion.UpdatedAt)
 }
+
+func listPublishedCanvasVersions(ctx context.Context, canvasUUID uuid.UUID, limit int, beforeTime *time.Time) (versions []models.CanvasVersion, count int64, err error) {
+	ctx, done := telemetry.Span(ctx, "canvases.list_published_versions")
+	defer done(&err)
+
+	err = database.DB(ctx).Transaction(func(tx *gorm.DB) error {
+		var txErr error
+		versions, txErr = models.ListPublishedCanvasVersionsInTransaction(tx, canvasUUID, limit, beforeTime)
+		if txErr != nil {
+			return txErr
+		}
+
+		count, txErr = models.CountPublishedCanvasVersionsInTransaction(tx, canvasUUID)
+		return txErr
+	})
+
+	return versions, count, err
+}
+
+func listDraftCanvasVersionsInTransaction(ctx context.Context, canvasID, ownerID uuid.UUID, limit int, beforeTime *time.Time) (versions []models.CanvasVersion, count int64, err error) {
+	ctx, done := telemetry.Span(ctx, "canvases.list_draft_versions")
+	defer done(&err)
+
+	err = database.DB(ctx).Transaction(func(tx *gorm.DB) error {
+		var txErr error
+		versions, txErr = models.ListDraftBranchesForCanvasInTransaction(tx, canvasID, ownerID, limit, beforeTime)
+		if txErr != nil {
+			return txErr
+		}
+
+		count, txErr = models.CountDraftBranchesForCanvasInTransaction(tx, canvasID, ownerID)
+		return txErr
+	})
+
+	return versions, count, err
+}