feat: more spans on scheduling

hatchet-dev · Aug 21, 2024 · 0fd27f1 · 0fd27f1
1 parent 326eeb3
commit 0fd27f1
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 0 deletions.
diff --git a/pkg/repository/prisma/step_run.go b/pkg/repository/prisma/step_run.go
@@ -849,6 +849,9 @@ type debugInfo struct {
 }
 
 func (s *stepRunEngineRepository) QueueStepRuns(ctx context.Context, tenantId string) (repository.QueueStepRunsResult, error) {
+	ctx, span := telemetry.NewSpan(ctx, "queue-step-runs-database")
+	defer span.End()
+
 	startedAt := time.Now().UTC()
 
 	emptyRes := repository.QueueStepRunsResult{
@@ -1092,6 +1095,7 @@ func (s *stepRunEngineRepository) QueueStepRuns(ctx context.Context, tenantId st
 	}
 
 	plan, err := scheduling.GeneratePlan(
+		ctx,
 		slots,
 		uniqueActionsArr,
 		queueItems,

diff --git a/pkg/scheduling/scheduling.go b/pkg/scheduling/scheduling.go
@@ -1,10 +1,12 @@
 package scheduling
 
 import (
+	"context"
 	"time"
 
 	"github.com/jackc/pgx/v5/pgtype"
 
+	"github.com/hatchet-dev/hatchet/internal/telemetry"
 	"github.com/hatchet-dev/hatchet/pkg/repository"
 	"github.com/hatchet-dev/hatchet/pkg/repository/prisma/dbsqlc"
 	"github.com/hatchet-dev/hatchet/pkg/repository/prisma/sqlchelpers"
@@ -17,6 +19,7 @@ type QueueItemWithOrder struct {
 }
 
 func GeneratePlan(
+	ctx context.Context,
 	slots []*dbsqlc.ListSemaphoreSlotsToAssignRow,
 	uniqueActionsArr []string,
 	queueItems []*QueueItemWithOrder,
@@ -25,6 +28,8 @@ func GeneratePlan(
 	workerLabels map[string][]*dbsqlc.GetWorkerLabelsRow,
 	stepDesiredLabels map[string][]*dbsqlc.GetDesiredLabelsRow,
 ) (SchedulePlan, error) {
+	_, span := telemetry.NewSpan(ctx, "generate-scheduling-plan")
+	defer span.End()
 
 	plan := SchedulePlan{
 		StepRunIds:             make([]pgtype.UUID, 0),

diff --git a/pkg/scheduling/scheduling_test.go b/pkg/scheduling/scheduling_test.go
@@ -1,6 +1,7 @@
 package scheduling
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"os"
@@ -258,6 +259,7 @@ func TestGeneratePlan(t *testing.T) {
 			}
 
 			got, err := GeneratePlan(
+				context.Background(),
 				fixtureData.Slots,
 				fixtureData.UniqueActionsArr,
 				fixtureData.QueueItems,
@@ -284,6 +286,7 @@ func BenchmarkGeneratePlan(b *testing.B) {
 
 	for i := 0; i < b.N; i++ {
 		_, _ = GeneratePlan(
+			context.Background(),
 			fixtureData.Slots,
 			fixtureData.UniqueActionsArr,
 			fixtureData.QueueItems,