Skip to content

Commit f13bf87

Browse files
igor-kupczynskiigor-kupczynski
authored
fix(api): cap tasks per workflow-run list read that includes payloads (#5205)
Payloads for a whole list page are loaded and serialized in memory, and a DAG run's child tasks are not bounded by the requested limit. Reject reads that would load payloads for more than 1000 tasks with a 400, and read DAG payloads once per DAG instead of once per finished child row.
1 parent 9d59e0d commit f13bf87

5 files changed

Lines changed: 212 additions & 14 deletions

File tree

‎api/v1/server/handlers/v1/workflow-runs/list.go‎

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package workflowruns
22

33
import (
44
"context"
5+
"errors"
56
"strings"
67
"time"
78

@@ -10,6 +11,7 @@ import (
1011

1112
"github.com/hatchet-dev/hatchet/api/v1/server/authz"
1213
v1handlers "github.com/hatchet-dev/hatchet/api/v1/server/handlers/v1"
14+
"github.com/hatchet-dev/hatchet/api/v1/server/oas/apierrors"
1315
"github.com/hatchet-dev/hatchet/api/v1/server/oas/gen"
1416
"github.com/hatchet-dev/hatchet/pkg/analytics"
1517
v1 "github.com/hatchet-dev/hatchet/pkg/repository"
@@ -231,7 +233,7 @@ func (t *V1WorkflowRunsService) shouldUseGinIndexForAdditionalMetadata(ctx conte
231233
return enabled || len(*additionalMetadata) == 1, nil
232234
}
233235

234-
func (t *V1WorkflowRunsService) WithDags(ctx context.Context, request gen.V1WorkflowRunListRequestObject, tenantId uuid.UUID, useGinIndex bool, canViewPayloads bool) (gen.V1WorkflowRunListResponseObject, error) {
236+
func (t *V1WorkflowRunsService) WithDags(ctx context.Context, request gen.V1WorkflowRunListRequestObject, tenantId uuid.UUID, useGinIndex bool, includePayloads bool, canViewPayloads bool) (gen.V1WorkflowRunListResponseObject, error) {
235237
ctx, span := telemetry.NewSpan(ctx, "v1-workflow-runs-list-with-dags-tasks")
236238
defer span.End()
237239

@@ -248,11 +250,6 @@ func (t *V1WorkflowRunsService) WithDags(ctx context.Context, request gen.V1Work
248250
offset = *request.Params.Offset
249251
}
250252

251-
includePayloads := false
252-
if request.Params.IncludePayloads != nil {
253-
includePayloads = *request.Params.IncludePayloads
254-
}
255-
256253
opts := listWorkflowRunOptsFromFilters(workflowRunFiltersFromListParams(request.Params), useGinIndex)
257254
opts.Limit = limit
258255
opts.Offset = offset
@@ -337,7 +334,7 @@ func (t *V1WorkflowRunsService) WithDags(ctx context.Context, request gen.V1Work
337334
), nil
338335
}
339336

340-
func (t *V1WorkflowRunsService) OnlyTasks(ctx context.Context, request gen.V1WorkflowRunListRequestObject, tenantId uuid.UUID, useGinIndex bool, canViewPayloads bool) (gen.V1WorkflowRunListResponseObject, error) {
337+
func (t *V1WorkflowRunsService) OnlyTasks(ctx context.Context, request gen.V1WorkflowRunListRequestObject, tenantId uuid.UUID, useGinIndex bool, includePayloads bool, canViewPayloads bool) (gen.V1WorkflowRunListResponseObject, error) {
341338
ctx, span := telemetry.NewSpan(ctx, "v1-workflow-runs-list-only-tasks")
342339
defer span.End()
343340

@@ -354,11 +351,6 @@ func (t *V1WorkflowRunsService) OnlyTasks(ctx context.Context, request gen.V1Wor
354351
offset = *request.Params.Offset
355352
}
356353

357-
includePayloads := false
358-
if request.Params.IncludePayloads != nil {
359-
includePayloads = *request.Params.IncludePayloads
360-
}
361-
362354
opts := listTaskRunOptsFromFilters(workflowRunFiltersFromListParams(request.Params), useGinIndex)
363355
opts.Limit = limit
364356
opts.Offset = offset
@@ -426,12 +418,21 @@ func (t *V1WorkflowRunsService) V1WorkflowRunList(ctx echo.Context, request gen.
426418
}
427419

428420
canViewPayloads := authz.CanViewPayloads(ctx)
421+
includePayloads := canViewPayloads && request.Params.IncludePayloads != nil && *request.Params.IncludePayloads
422+
423+
var res gen.V1WorkflowRunListResponseObject
429424

430425
if request.Params.OnlyTasks {
431-
return t.OnlyTasks(spanContext, request, tenantId, useGinIndex, canViewPayloads)
426+
res, err = t.OnlyTasks(spanContext, request, tenantId, useGinIndex, includePayloads, canViewPayloads)
427+
} else {
428+
res, err = t.WithDags(spanContext, request, tenantId, useGinIndex, includePayloads, canViewPayloads)
429+
}
430+
431+
if errors.Is(err, v1.ErrTooManyTasksWithPayloads) {
432+
return gen.V1WorkflowRunList400JSONResponse(apierrors.NewAPIErrors(err.Error())), nil
432433
}
433434

434-
return t.WithDags(spanContext, request, tenantId, useGinIndex, canViewPayloads)
435+
return res, err
435436
}
436437

437438
func (t *V1WorkflowRunsService) V1WorkflowRunCountGet(ctx echo.Context, request gen.V1WorkflowRunCountGetRequestObject) (gen.V1WorkflowRunCountGetResponseObject, error) {

‎pkg/repository/olap.go‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,25 @@ const (
4343
AdditionalMetadataOperatorAnd AdditionalMetadataOperator = "AND"
4444
)
4545

46+
// MaxTasksWithPayloads caps how many runs, and separately how many DAG child tasks, a list read
47+
// loads payloads for. Payloads for a whole page are held in memory and serialized at once. The
48+
// check runs on the rows a query returned rather than on the requested limit, so a large limit
49+
// over a sparse window still works.
50+
const MaxTasksWithPayloads = 1000
51+
52+
var ErrTooManyTasksWithPayloads = fmt.Errorf(
53+
"cannot include payloads for a page with more than %d runs, or whose DAG runs have more than %d tasks in total; lower the limit or set include_payloads=false",
54+
MaxTasksWithPayloads, MaxTasksWithPayloads,
55+
)
56+
57+
func checkTasksWithPayloads(includePayloads bool, numTasks int) error {
58+
if includePayloads && numTasks > MaxTasksWithPayloads {
59+
return ErrTooManyTasksWithPayloads
60+
}
61+
62+
return nil
63+
}
64+
4665
type ListTaskRunOpts struct {
4766
CreatedAfter time.Time
4867

@@ -898,6 +917,10 @@ func (r *OLAPRepositoryImpl) ListTasks(ctx context.Context, tenantId uuid.UUID,
898917
return nil, 0, err
899918
}
900919

920+
if err := checkTasksWithPayloads(opts.IncludePayloads, len(rows)); err != nil {
921+
return nil, 0, err
922+
}
923+
901924
if opts.IncludeTotalCount {
902925
count, err = r.queries.CountTasks(ctx, tx, countTasksParamsFromListTasksParams(params))
903926

@@ -1004,6 +1027,11 @@ func (r *OLAPRepositoryImpl) ListTasksByDAGId(ctx context.Context, tenantId uuid
10041027
return nil, taskIdToDagExternalId, err
10051028
}
10061029

1030+
// The page limit bounds the DAG runs, not their child tasks.
1031+
if err := checkTasksWithPayloads(includePayloads, len(tasks)); err != nil {
1032+
return nil, taskIdToDagExternalId, err
1033+
}
1034+
10071035
idsInsertedAts := make([]IdInsertedAt, 0, len(tasks))
10081036

10091037
for _, row := range tasks {
@@ -1376,6 +1404,10 @@ func (r *OLAPRepositoryImpl) ListWorkflowRuns(ctx context.Context, tenantId uuid
13761404
return nil, 0, err
13771405
}
13781406

1407+
if err := checkTasksWithPayloads(opts.IncludePayloads, len(workflowRunIds)); err != nil {
1408+
return nil, 0, err
1409+
}
1410+
13791411
if opts.IncludeTotalCount {
13801412
count, err = r.queries.CountWorkflowRuns(ctx, tx, countWorkflowRunsParamsFromFetchWorkflowRunIdsParams(params))
13811413

‎pkg/repository/olap_list_workflow_runs_test.go‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
//go:build !e2e && !load && !rampup && !integration
2+
3+
package repository
4+
5+
import (
6+
"context"
7+
"testing"
8+
"time"
9+
10+
"github.com/google/uuid"
11+
"github.com/jackc/pgx/v5/pgtype"
12+
"github.com/jackc/pgx/v5/pgxpool"
13+
"github.com/stretchr/testify/require"
14+
15+
"github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1"
16+
)
17+
18+
type finishedDAGFixture struct {
19+
tenantId uuid.UUID
20+
dagId int64
21+
insertedAt pgtype.Timestamptz
22+
lastFinishedEventId uuid.UUID
23+
}
24+
25+
func setupOLAPRepository(t *testing.T) (context.Context, *pgxpool.Pool, *OLAPRepositoryImpl) {
26+
t.Helper()
27+
28+
basePool, cleanup := setupPostgresWithMigration(t)
29+
t.Cleanup(cleanup)
30+
31+
pool := createEnumAwarePool(t, basePool)
32+
repo := createOLAPRepositoryWithPayloadStore(t, pool)
33+
34+
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
35+
t.Cleanup(cancel)
36+
37+
require.NoError(t, repo.UpdateTablePartitions(ctx))
38+
39+
return ctx, pool, repo
40+
}
41+
42+
// seedDAGWithFinishedChildren writes one DAG run whose children each finish in their own write, so
43+
// the FINISHED events get distinct inserted_at values.
44+
func seedDAGWithFinishedChildren(t *testing.T, ctx context.Context, repo *OLAPRepositoryImpl, numChildren int) finishedDAGFixture {
45+
t.Helper()
46+
47+
f := finishedDAGFixture{
48+
tenantId: uuid.New(),
49+
dagId: 1,
50+
insertedAt: pgtype.Timestamptz{Time: time.Now().UTC().Truncate(time.Microsecond), Valid: true},
51+
}
52+
53+
workflowId := uuid.New()
54+
workflowVersionId := uuid.New()
55+
dagExternalId := uuid.New()
56+
57+
locksNotAcquired, err := repo.CreateDAGs(ctx, f.tenantId, []*DAGWithData{{
58+
V1Dag: &sqlcv1.V1Dag{
59+
ID: f.dagId,
60+
InsertedAt: f.insertedAt,
61+
TenantID: f.tenantId,
62+
ExternalID: dagExternalId,
63+
DisplayName: "finished-children-dag",
64+
WorkflowID: workflowId,
65+
WorkflowVersionID: workflowVersionId,
66+
},
67+
Input: []byte(`{}`),
68+
AdditionalMetadata: []byte(`{}`),
69+
}})
70+
require.NoError(t, err)
71+
require.Empty(t, locksNotAcquired)
72+
73+
children := make([]*V1TaskWithPayload, 0, numChildren)
74+
75+
for i := range numChildren {
76+
children = append(children, &V1TaskWithPayload{
77+
V1Task: &sqlcv1.V1Task{
78+
ID: f.dagId + int64(i) + 1,
79+
InsertedAt: f.insertedAt,
80+
TenantID: f.tenantId,
81+
Queue: "default",
82+
ActionID: "test:finished-children-dag-child",
83+
StepID: uuid.New(),
84+
WorkflowID: workflowId,
85+
WorkflowVersionID: workflowVersionId,
86+
WorkflowRunID: dagExternalId,
87+
ScheduleTimeout: "5m",
88+
Sticky: sqlcv1.V1StickyStrategyNONE,
89+
ExternalID: uuid.New(),
90+
DisplayName: "finished-children-dag-child",
91+
Input: []byte(`{}`),
92+
AdditionalMetadata: []byte(`{}`),
93+
DagID: pgtype.Int8{Int64: f.dagId, Valid: true},
94+
DagInsertedAt: f.insertedAt,
95+
},
96+
Payload: []byte(`{}`),
97+
})
98+
}
99+
100+
_, locksNotAcquired, err = repo.CreateTasks(ctx, f.tenantId, children)
101+
require.NoError(t, err)
102+
require.Empty(t, locksNotAcquired)
103+
104+
for _, child := range children {
105+
event := sqlcv1.CreateTaskEventsOLAPParams{
106+
TenantID: f.tenantId,
107+
TaskID: child.ID,
108+
TaskInsertedAt: f.insertedAt,
109+
EventType: sqlcv1.V1EventTypeOlapFINISHED,
110+
WorkflowID: workflowId,
111+
EventTimestamp: pgtype.Timestamptz{Time: time.Now().UTC(), Valid: true},
112+
ReadableStatus: sqlcv1.V1ReadableStatusOlapCOMPLETED,
113+
Output: []byte(`{}`),
114+
ExternalID: uuid.New(),
115+
}
116+
f.lastFinishedEventId = event.ExternalID
117+
118+
_, locksNotAcquired, err = repo.CreateTaskEvents(ctx, f.tenantId, []sqlcv1.CreateTaskEventsOLAPParams{event}, map[uuid.UUID]uuid.UUID{event.ExternalID: dagExternalId}, nil, nil)
119+
require.NoError(t, err)
120+
require.Empty(t, locksNotAcquired)
121+
}
122+
123+
return f
124+
}
125+
126+
func TestPopulateDAGMetadata_OneRowPerDAGWithLatestChildOutput(t *testing.T) {
127+
ctx, pool, repo := setupOLAPRepository(t)
128+
f := seedDAGWithFinishedChildren(t, ctx, repo, 2)
129+
130+
rows, err := repo.queries.PopulateDAGMetadata(ctx, pool, sqlcv1.PopulateDAGMetadataParams{
131+
Ids: []int64{f.dagId},
132+
Insertedats: []pgtype.Timestamptz{f.insertedAt},
133+
Tenantid: f.tenantId,
134+
Includepayloads: true,
135+
})
136+
require.NoError(t, err)
137+
require.Len(t, rows, 1)
138+
require.NotNil(t, rows[0].OutputEventExternalID)
139+
require.Equal(t, f.lastFinishedEventId, *rows[0].OutputEventExternalID)
140+
}
141+
142+
func TestListWorkflowRuns_PayloadCapCountsReturnedRowsNotLimit(t *testing.T) {
143+
ctx, _, repo := setupOLAPRepository(t)
144+
f := seedDAGWithFinishedChildren(t, ctx, repo, 2)
145+
146+
runs, _, err := repo.ListWorkflowRuns(ctx, f.tenantId, ListWorkflowRunOpts{
147+
CreatedAfter: f.insertedAt.Time.Add(-time.Hour),
148+
Limit: MaxTasksWithPayloads * 5,
149+
IncludePayloads: true,
150+
})
151+
require.NoError(t, err)
152+
require.Len(t, runs, 1)
153+
}
154+
155+
func TestCheckTasksWithPayloads(t *testing.T) {
156+
require.NoError(t, checkTasksWithPayloads(true, MaxTasksWithPayloads))
157+
require.ErrorIs(t, checkTasksWithPayloads(true, MaxTasksWithPayloads+1), ErrTooManyTasksWithPayloads)
158+
require.NoError(t, checkTasksWithPayloads(false, MaxTasksWithPayloads+1))
159+
}

‎pkg/repository/sqlcv1/olap.sql‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1534,6 +1534,7 @@ WITH input AS (
15341534
e.run_id, e.retry_count DESC
15351535
), task_output AS (
15361536
SELECT
1537+
DISTINCT ON (run_id)
15371538
run_id,
15381539
output,
15391540
external_id,
@@ -1542,6 +1543,8 @@ WITH input AS (
15421543
relevant_events
15431544
WHERE
15441545
event_type = 'FINISHED'
1546+
ORDER BY
1547+
run_id, inserted_at DESC
15451548
)
15461549

15471550
SELECT

‎pkg/repository/sqlcv1/olap.sql.go‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)