refactor(engine): simplify task data shape (#5198)

* refactor: simplify the V1StepRunData and read error messages out of the task events

* feat: add new columns to db

* chore: gen

* feat: wiring

* refactor: simplify a bit

* fix: migration version

* fix: handle replay bug

* fix: add an env var to gate

* fix: simplify by adding a payload version

* fix: simplify more

* fix: simplify yet more, remove some tests

* fix: make filter payloads compatible across versions

* chore: shrink diff more

Author: mrkaye97

cmd/hatchet-migrate/migrate/migrations/20261014180746_v1_0_174.sql

@@ -0,0 +1,17 @@
+-- +goose Up
+-- +goose StatementBegin
+ALTER TABLE v1_task
+    ADD COLUMN triggers JSONB,
+    ADD COLUMN dag_parent_task_run_ids UUID[],
+    ADD COLUMN child_task_desired_worker_labels JSONB,
+    ADD COLUMN task_input_payload_version SMALLINT NOT NULL DEFAULT 1;
+-- +goose StatementEnd
+
+-- +goose Down
+-- +goose StatementBegin
+ALTER TABLE v1_task
+    DROP COLUMN triggers,
+    DROP COLUMN dag_parent_task_run_ids,
+    DROP COLUMN child_task_desired_worker_labels,
+    DROP COLUMN task_input_payload_version;
+-- +goose StatementEnd

internal/services/dispatcher/dispatcher.go

@@ -810,13 +810,11 @@ func (d *DispatcherImpl) populateTaskData(
 			input = task.Input
 		}
 
-		currInput := &v1.V1StepRunData{}
+		currInput, err := v1.NewV1StepRunDataFromTaskAndInputPayload(task, input)
 
-		if input != nil {
-			if err := json.Unmarshal(input, currInput); err != nil {
-				d.l.Warn().Ctx(ctx).Err(err).Msg("failed to unmarshal input")
-				continue
-			}
+		if err != nil {
+			d.l.Warn().Ctx(ctx).Err(err).Msg("failed to unmarshal input")
+			continue
 		}
 
 		if len(currInput.DagParentTaskRunIds) > 0 {
@@ -826,24 +824,13 @@ func (d *DispatcherImpl) populateTaskData(
 				dagParentExternalIdSet[parentExternalId] = struct{}{}
 			}
 		} else if parentData, ok := parentDataMap[task.ID]; ok {
-			readableIdToData := make(map[string]map[string]interface{})
-
-			for _, outputEvent := range parentData {
-				outputMap := make(map[string]interface{})
-
-				if len(outputEvent.Output) > 0 {
-					if err := json.Unmarshal(outputEvent.Output, &outputMap); err != nil {
-						d.l.Warn().Ctx(ctx).Err(err).Msg("failed to unmarshal output")
-						continue
-					}
-				}
-
-				readableIdToData[outputEvent.StepReadableID] = outputMap
-			}
+			parentOutputs, stepRunErrors := splitParentOutputsAndStepRunErrors(ctx, d.l, task.ID, parentData)
 
-			currInput.Parents = readableIdToData
-			inputs[payloadKey] = currInput.Bytes()
+			currInput.Parents = parentOutputs
+			currInput.StepRunErrors = stepRunErrors
 		}
+
+		inputs[payloadKey] = currInput.Bytes()
 	}
 
 	if len(dagChildInputs) > 0 {
@@ -921,6 +908,39 @@ func (d *DispatcherImpl) populateTaskData(
 	return taskIdToData, nil
 }
 
+func splitParentOutputsAndStepRunErrors(
+	ctx context.Context,
+	l *zerolog.Logger,
+	taskId int64,
+	terminalEvents []*v1.TaskOutputEvent,
+) (map[string]map[string]interface{}, map[string]string) {
+	parentOutputs := make(map[string]map[string]interface{})
+	stepRunErrors := make(map[string]string)
+
+	for _, event := range terminalEvents {
+		if event.IsFailed() {
+			if event.TaskId != taskId {
+				stepRunErrors[event.StepReadableID] = event.ErrorMessage
+			}
+
+			continue
+		}
+
+		outputMap := make(map[string]interface{})
+
+		if len(event.Output) > 0 {
+			if err := json.Unmarshal(event.Output, &outputMap); err != nil {
+				l.Warn().Ctx(ctx).Err(err).Msg("failed to unmarshal output")
+				continue
+			}
+		}
+
+		parentOutputs[event.StepReadableID] = outputMap
+	}
+
+	return parentOutputs, stepRunErrors
+}
+
 // dagChildTaskInput is a durable DAG-operator child task awaiting parent output resolution:
 // its own payload key (to write the resolved input back into `inputs`) plus its unmarshaled
 // input (which carries the list of parent external IDs it needs outputs for).

pkg/config/loader/loader.go

@@ -435,6 +435,7 @@ func (c *ConfigLoader) InitDataLayer() (res *database.Layer, err error) {
 		ExternalCutoverNumConcurrentOffloads: scf.PayloadStore.ExternalCutoverNumConcurrentOffloads,
 		InlineStoreTTL:                       &inlineStoreTTL,
 		EnableWindowSizeOptimization:         scf.PayloadStore.EnableWindowSizeOptimization,
+		WrapTaskInputInStepRunData:           scf.PayloadStore.WrapTaskInputInStepRunData,
 	}
 
 	statusUpdateOpts := repov1.StatusUpdateBatchSizeLimits{

pkg/config/server/server.go

@@ -826,6 +826,12 @@ type PayloadStoreConfig struct {
 	ExternalCutoverNumConcurrentOffloads int32         `mapstructure:"externalCutoverNumConcurrentOffloads" json:"externalCutoverNumConcurrentOffloads,omitempty" default:"10"`
 	InlineStoreTTLDays                   int32         `mapstructure:"inlineStoreTTLDays" json:"inlineStoreTTLDays,omitempty" default:"2"`
 	EnableWindowSizeOptimization         bool          `mapstructure:"enableWindowSizeOptimization" json:"enableWindowSizeOptimization,omitempty" default:"true"`
+	// Writes the task input to the payload store wrapped in the full step run data, which engines before the
+	// `triggers`, `dag_parent_task_run_ids` and `child_task_desired_worker_labels` columns on v1_task can read.
+	// Every engine from this version on reads both shapes. Turn off only once every engine is on this version
+	// or later, because an older engine cannot read a bare input and turning it back on does not repair rows
+	// already written bare.
+	WrapTaskInputInStepRunData bool `mapstructure:"wrapTaskInputInStepRunData" json:"wrapTaskInputInStepRunData,omitempty" default:"true"`
 }
 
 func (c *ServerConfig) HasService(name string) bool {
@@ -1120,6 +1126,7 @@ func BindAllEnv(v *viper.Viper) {
 	_ = v.BindEnv("payloadStore.externalCutoverNumConcurrentOffloads", "SERVER_PAYLOAD_STORE_EXTERNAL_CUTOVER_NUM_CONCURRENT_OFFLOADS")
 	_ = v.BindEnv("payloadStore.inlineStoreTTLDays", "SERVER_PAYLOAD_STORE_INLINE_STORE_TTL_DAYS")
 	_ = v.BindEnv("payloadStore.enableWindowSizeOptimization", "SERVER_PAYLOAD_STORE_ENABLE_WINDOW_SIZE_OPTIMIZATION")
+	_ = v.BindEnv("payloadStore.wrapTaskInputInStepRunData", "SERVER_PAYLOAD_STORE_WRAP_TASK_INPUT_IN_STEP_RUN_DATA")
 
 	// cron operations options
 	_ = v.BindEnv("cronOperations.taskAnalyzeCronInterval", "SERVER_CRON_OPERATIONS_TASK_ANALYZE_CRON_INTERVAL")

pkg/repository/input.go

@@ -2,6 +2,7 @@ package repository
 
 import (
 	"encoding/json"
+	"fmt"
 
 	"github.com/google/uuid"
 
@@ -35,12 +36,35 @@ func (s *sharedRepository) DesiredWorkerId(t *TaskInput) *uuid.UUID {
 	return nil
 }
 
-func (s *sharedRepository) newTaskInputFromExistingBytes(inputBytes []byte) *TaskInput {
+func (s *sharedRepository) newTaskInputFromTaskInputPayload(version taskInputPayloadVersion, taskInputPayload []byte) *TaskInput {
 	i := &TaskInput{}
 
-	err := json.Unmarshal(inputBytes, i)
+	var userInput []byte
 
-	if err != nil {
+	switch version {
+	case taskInputPayloadVersionStepRunData:
+		stepRunData, err := stepRunDataFromVersion1TaskInputPayload(taskInputPayload)
+
+		if err != nil {
+			s.l.Error().Err(err).Msg("failed to unmarshal step run data")
+			return i
+		}
+
+		userInput = stepRunData.Input
+		i.DagParentTaskRunIds = stepRunData.DagParentTaskRunIds
+		i.DesiredWorkerLabels = stepRunData.DesiredWorkerLabels
+	case taskInputPayloadVersionUserInputOnly:
+		userInput = taskInputPayload
+	default:
+		s.l.Error().Msgf("unknown task input payload version %d", version)
+		return i
+	}
+
+	if len(userInput) == 0 {
+		return i
+	}
+
+	if err := json.Unmarshal(userInput, &i.Input); err != nil {
 		s.l.Error().Err(err).Msg("failed to unmarshal input bytes")
 	}
 
@@ -126,9 +150,14 @@ func (s *sharedRepository) ToV1StepRunData(t *TaskInput) *V1StepRunData {
 
 	triggers["filter_payload"] = t.FilterPayload
 
+	input, err := json.Marshal(t.Input)
+
+	if err != nil {
+		s.l.Warn().Err(err).Msg("failed to marshal input")
+	}
+
 	return &V1StepRunData{
-		Input:               t.Input,
-		TriggeredBy:         "manual",
+		Input:               input,
 		Parents:             parents,
 		Triggers:            triggers,
 		StepRunErrors:       stepRunErrors,
@@ -137,36 +166,104 @@ func (s *sharedRepository) ToV1StepRunData(t *TaskInput) *V1StepRunData {
 	}
 }
 
-func (r *sharedRepository) V1StepRunDataFromBytes(b []byte) (*V1StepRunData, error) {
-	if len(b) == 0 {
-		return nil, nil
+type taskInputPayloadVersion int16
+
+const (
+	taskInputPayloadVersionStepRunData   taskInputPayloadVersion = 1
+	taskInputPayloadVersionUserInputOnly taskInputPayloadVersion = 2
+)
+
+func (s *sharedRepository) taskInputPayloadVersionToWrite() taskInputPayloadVersion {
+	if s.wrapTaskInputInStepRunData {
+		return taskInputPayloadVersionStepRunData
 	}
 
-	out := &V1StepRunData{}
+	return taskInputPayloadVersionUserInputOnly
+}
 
-	err := json.Unmarshal(b, out)
+func stepRunDataFromVersion1TaskInputPayload(taskInputPayload []byte) (*V1StepRunData, error) {
+	stepRunData := &V1StepRunData{}
 
-	if err != nil {
+	if len(taskInputPayload) == 0 {
+		return stepRunData, nil
+	}
+
+	if err := json.Unmarshal(taskInputPayload, stepRunData); err != nil {
 		return nil, err
 	}
 
-	return out, nil
+	return stepRunData, nil
 }
 
-type V1StepRunData struct {
-	Input       map[string]interface{}            `json:"input"`
-	TriggeredBy string                            `json:"triggered_by"`
-	Parents     map[string]map[string]interface{} `json:"parents"`
+func userInputFromTaskInputPayload(version taskInputPayloadVersion, taskInputPayload []byte) ([]byte, error) {
+	switch version {
+	case taskInputPayloadVersionStepRunData:
+		stepRunData, err := stepRunDataFromVersion1TaskInputPayload(taskInputPayload)
 
-	Triggers map[string]map[string]interface{} `json:"triggers"`
+		if err != nil {
+			return nil, fmt.Errorf("failed to unmarshal step run data: %w", err)
+		}
+
+		return stepRunData.Input, nil
+	case taskInputPayloadVersionUserInputOnly:
+		return taskInputPayload, nil
+	default:
+		return nil, fmt.Errorf("unknown task input payload version %d", version)
+	}
+}
+
+func NewV1StepRunDataFromTaskAndInputPayload(task *sqlcv1.V1Task, taskInputPayload []byte) (*V1StepRunData, error) {
+	switch version := taskInputPayloadVersion(task.TaskInputPayloadVersion); version {
+	case taskInputPayloadVersionStepRunData:
+		return stepRunDataFromVersion1TaskInputPayload(taskInputPayload)
+	case taskInputPayloadVersionUserInputOnly:
+		return newV1StepRunDataFromTaskAndUserInput(task, taskInputPayload)
+	default:
+		return nil, fmt.Errorf("unknown task input payload version %d", version)
+	}
+}
+
+// newV1StepRunDataFromTaskAndUserInput leaves Parents empty and StepRunErrors unset. The dispatcher fills both in
+// from the task events of the task's parents.
+func newV1StepRunDataFromTaskAndUserInput(task *sqlcv1.V1Task, userInput []byte) (*V1StepRunData, error) {
+	triggers := map[string]map[string]interface{}{"filter_payload": nil}
+
+	if len(task.Triggers) > 0 {
+		if err := json.Unmarshal(task.Triggers, &triggers); err != nil {
+			return nil, fmt.Errorf("failed to unmarshal triggers: %w", err)
+		}
+	}
+
+	var childTaskDesiredWorkerLabels []*sqlcv1.GetDesiredLabelsRow
+
+	if len(task.ChildTaskDesiredWorkerLabels) > 0 {
+		if err := json.Unmarshal(task.ChildTaskDesiredWorkerLabels, &childTaskDesiredWorkerLabels); err != nil {
+			return nil, fmt.Errorf("failed to unmarshal child task desired worker labels: %w", err)
+		}
+	}
+
+	input := json.RawMessage(userInput)
 
-	// custom-set user data for the step
-	UserData map[string]interface{} `json:"user_data"`
+	if len(input) == 0 {
+		input = json.RawMessage("{}")
+	}
+
+	return &V1StepRunData{
+		Input:               input,
+		Parents:             make(map[string]map[string]interface{}),
+		Triggers:            triggers,
+		DagParentTaskRunIds: task.DagParentTaskRunIds,
+		DesiredWorkerLabels: childTaskDesiredWorkerLabels,
+	}, nil
+}
 
-	// overrides set from the playground
-	Overrides map[string]interface{} `json:"overrides"`
+type V1StepRunData struct {
+	Input    json.RawMessage                   `json:"input"`
+	Parents  map[string]map[string]interface{} `json:"parents"`
+	Triggers map[string]map[string]interface{} `json:"triggers"`
 
-	// errors in upstream steps (only used in on-failure step)
+	// errors in upstream steps. Set at task creation for engines that read the wrapped payload, and overwritten
+	// from the parents' task events when the task is dispatched.
 	StepRunErrors map[string]string `json:"step_run_errors,omitempty"`
 
 	// task run external IDs of parent tasks in a durable DAG orchestration
@@ -177,25 +274,67 @@ type V1StepRunData struct {
 	DesiredWorkerLabels []*sqlcv1.GetDesiredLabelsRow `json:"desired_worker_labels,omitempty"`
 }
 
-func (v1 *V1StepRunData) Bytes() []byte {
-	out, err := json.Marshal(v1)
+func (v1 *V1StepRunData) taskInputPayload(version taskInputPayloadVersion) []byte {
+	switch version {
+	case taskInputPayloadVersionStepRunData:
+		return v1.Bytes()
+	case taskInputPayloadVersionUserInputOnly:
+		return v1.userInput()
+	default:
+		panic(fmt.Sprintf("unknown task input payload version %d", version))
+	}
+}
 
-	if err != nil {
-		return []byte("{}")
+// userInput is null for a task created without an input, which matches what Bytes returns for it.
+func (v1 *V1StepRunData) userInput() []byte {
+	if v1 == nil {
+		return []byte("null")
 	}
 
-	return out
+	return v1.Input
+}
+
+type stepRunDataTaskColumns struct {
+	triggers                     []byte
+	dagParentTaskRunIds          []byte
+	childTaskDesiredWorkerLabels []byte
 }
 
-func (v1 *V1StepRunData) InputBytes() []byte {
+func (v1 *V1StepRunData) taskColumns() (stepRunDataTaskColumns, error) {
+	columns := stepRunDataTaskColumns{}
+
 	if v1 == nil {
-		return nil
+		return columns, nil
+	}
+
+	var err error
+
+	if len(v1.Triggers) > 0 {
+		if columns.triggers, err = json.Marshal(v1.Triggers); err != nil {
+			return columns, fmt.Errorf("failed to marshal triggers: %w", err)
+		}
 	}
 
-	out, err := json.Marshal(v1.Input)
+	if len(v1.DagParentTaskRunIds) > 0 {
+		if columns.dagParentTaskRunIds, err = json.Marshal(v1.DagParentTaskRunIds); err != nil {
+			return columns, fmt.Errorf("failed to marshal dag parent task run ids: %w", err)
+		}
+	}
+
+	if len(v1.DesiredWorkerLabels) > 0 {
+		if columns.childTaskDesiredWorkerLabels, err = json.Marshal(v1.DesiredWorkerLabels); err != nil {
+			return columns, fmt.Errorf("failed to marshal child task desired worker labels: %w", err)
+		}
+	}
+
+	return columns, nil
+}
+
+func (v1 *V1StepRunData) Bytes() []byte {
+	out, err := json.Marshal(v1)
 
 	if err != nil {
-		return nil
+		return []byte("{}")
 	}
 
 	return out