Skip to content

Commit 74d3f57

Browse files
author
bgagent
committed
feat(orchestration): admission queue with deferred pickup (#441)
When a task hits the per-user concurrency cap, park it in a new QUEUED state instead of failing it (the #331 mass-fail mode). A scheduled AdmissionQueuePickup Lambda drains the queue FIFO (by created_at) as slots free up, flipping QUEUED -> SUBMITTED and re-invoking the orchestrator, whose atomic admissionControl stays the single writer of the concurrency counter — a lost pickup race harmlessly re-queues without losing FIFO position. - task-status: new QUEUED state (pre-active — holds no slot), with SUBMITTED <-> QUEUED transitions plus QUEUED -> CANCELLED/FAILED - orchestrator: queueTask() replaces failTask on cap; queued_at is stamped once, admission_attempts increments per pass; Linear/Jira channel feedback now says 'queued' and fires only on first entry - reconcile-admission-queue: FIFO drain per user with read-only capacity pre-check, conditional flips (cancel-safe), 24h max-age backstop, and orchestrator re-invoke carrying a pickup nonce - reconcile-stranded-tasks: age by time-in-current-status so a task that waited in the queue longer than the stranded timeout is not killed right after pickup - API/CLI: GET /tasks/{id} computes read-time queue_position + estimated_wait_s (fail-open); bgagent status/detail render a 'Queue: position N (est. wait ~Xm)' line - cancel: QUEUED is cancellable everywhere (REST already generic; Slack cancel list extended); no concurrency release needed - idempotent replay returns the existing QUEUED task unchanged - tests: state-machine invariants, queueTask, pickup handler (incl. #331 fan-out-burst regression: burst above cap survives and drains FIFO with zero failures), get-task queue position, CLI rendering Closes #441
1 parent 7723355 commit 74d3f57

22 files changed

Lines changed: 1648 additions & 57 deletions
Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
/**
2+
* MIT No Attribution
3+
*
4+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
5+
*
6+
* Permission is hereby granted, free of charge, to any person obtaining a copy of
7+
* the Software without restriction, including without limitation the rights to
8+
* use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
9+
* the Software, and to permit persons to whom the Software is furnished to do so.
10+
*
11+
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
12+
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
13+
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
14+
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
15+
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
16+
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
17+
* SOFTWARE.
18+
*/
19+
20+
import * as path from 'path';
21+
import { Duration } from 'aws-cdk-lib';
22+
import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
23+
import * as events from 'aws-cdk-lib/aws-events';
24+
import * as targets from 'aws-cdk-lib/aws-events-targets';
25+
import * as iam from 'aws-cdk-lib/aws-iam';
26+
import { Architecture, Runtime } from 'aws-cdk-lib/aws-lambda';
27+
import * as lambda from 'aws-cdk-lib/aws-lambda-nodejs';
28+
import { NagSuppressions } from 'cdk-nag';
29+
import { Construct } from 'constructs';
30+
31+
/** Pickup Lambda timeout (minutes). */
32+
const PICKUP_TIMEOUT_MINUTES = 5;
33+
34+
/** Pickup Lambda memory (MB). */
35+
const PICKUP_MEMORY_MB = 256;
36+
37+
/**
38+
* Default pickup schedule interval (minutes). Short — queue latency is
39+
* user-visible wait time; 1 minute keeps the worst-case pickup delay
40+
* bounded while a QUEUED-empty cycle is a single cheap GSI query.
41+
*/
42+
const DEFAULT_SCHEDULE_MINUTES = 1;
43+
44+
/** Default max queue age before the backstop fails the task (seconds; 24h). */
45+
const DEFAULT_QUEUE_MAX_AGE_SECONDS = 86400;
46+
47+
/** Default task-record retention used for event TTL (days). */
48+
const DEFAULT_TASK_RETENTION_DAYS = 90;
49+
50+
/**
51+
* Properties for AdmissionQueuePickup construct.
52+
*/
53+
export interface AdmissionQueuePickupProps {
54+
/** TaskTable (StatusIndex GSI powers the QUEUED FIFO query). */
55+
readonly taskTable: dynamodb.ITable;
56+
57+
/** TaskEventsTable (handler writes queue_pickup / task_failed events). */
58+
readonly taskEventsTable: dynamodb.ITable;
59+
60+
/** UserConcurrencyTable (read-only capacity check per user). */
61+
readonly userConcurrencyTable: dynamodb.ITable;
62+
63+
/** ARN of the orchestrator Lambda alias to re-invoke on pickup. */
64+
readonly orchestratorFunctionArn: string;
65+
66+
/**
67+
* Maximum concurrent tasks per user — must match the orchestrator's
68+
* value so the capacity pre-check agrees with `admissionControl`.
69+
* @default 10
70+
*/
71+
readonly maxConcurrentTasksPerUser?: number;
72+
73+
/**
74+
* How often to drain the queue.
75+
* @default Duration.minutes(1)
76+
*/
77+
readonly schedule?: Duration;
78+
79+
/**
80+
* Max time a task may sit QUEUED before the backstop fails it (seconds).
81+
* @default 86400 (24 hours)
82+
*/
83+
readonly queueMaxAgeSeconds?: number;
84+
85+
/** Forwarded to the handler for event TTL. @default 90 */
86+
readonly taskRetentionDays?: number;
87+
}
88+
89+
/**
90+
* Scheduled Lambda that drains the admission queue (#441).
91+
*
92+
* Tasks that hit the per-user concurrency cap are parked in QUEUED by the
93+
* orchestrator instead of FAILED. This Lambda re-attempts admission in
94+
* FIFO order (StatusIndex GSI, ascending ``created_at``) as slots free
95+
* up: it flips QUEUED -> SUBMITTED and re-invokes the orchestrator, whose
96+
* atomic `admissionControl` remains the single writer of the concurrency
97+
* counter (a lost race simply re-queues the task, preserving position).
98+
*/
99+
export class AdmissionQueuePickup extends Construct {
100+
public readonly fn: lambda.NodejsFunction;
101+
102+
constructor(scope: Construct, id: string, props: AdmissionQueuePickupProps) {
103+
super(scope, id);
104+
105+
const handlersDir = path.join(__dirname, '..', 'handlers');
106+
107+
this.fn = new lambda.NodejsFunction(this, 'PickupFn', {
108+
entry: path.join(handlersDir, 'reconcile-admission-queue.ts'),
109+
handler: 'handler',
110+
runtime: Runtime.NODEJS_24_X,
111+
architecture: Architecture.ARM_64,
112+
timeout: Duration.minutes(PICKUP_TIMEOUT_MINUTES),
113+
memorySize: PICKUP_MEMORY_MB,
114+
environment: {
115+
TASK_TABLE_NAME: props.taskTable.tableName,
116+
TASK_EVENTS_TABLE_NAME: props.taskEventsTable.tableName,
117+
USER_CONCURRENCY_TABLE_NAME: props.userConcurrencyTable.tableName,
118+
ORCHESTRATOR_FUNCTION_ARN: props.orchestratorFunctionArn,
119+
MAX_CONCURRENT_TASKS_PER_USER: String(props.maxConcurrentTasksPerUser ?? 10),
120+
QUEUE_MAX_AGE_SECONDS: String(props.queueMaxAgeSeconds ?? DEFAULT_QUEUE_MAX_AGE_SECONDS),
121+
TASK_RETENTION_DAYS: String(props.taskRetentionDays ?? DEFAULT_TASK_RETENTION_DAYS),
122+
},
123+
bundling: {
124+
externalModules: ['@aws-sdk/*'],
125+
},
126+
});
127+
128+
// TaskTable: StatusIndex query + conditional QUEUED->SUBMITTED /
129+
// QUEUED->FAILED transitions.
130+
props.taskTable.grantReadWriteData(this.fn);
131+
// TaskEvents: queue_pickup / task_failed events.
132+
props.taskEventsTable.grantWriteData(this.fn);
133+
// Concurrency: READ-ONLY capacity pre-check — the orchestrator's
134+
// admissionControl is the only writer of the counter.
135+
props.userConcurrencyTable.grantReadData(this.fn);
136+
137+
// Re-invoke the orchestrator alias on pickup.
138+
this.fn.addToRolePolicy(new iam.PolicyStatement({
139+
actions: ['lambda:InvokeFunction'],
140+
resources: [props.orchestratorFunctionArn],
141+
}));
142+
143+
const schedule = props.schedule ?? Duration.minutes(DEFAULT_SCHEDULE_MINUTES);
144+
const rule = new events.Rule(this, 'PickupSchedule', {
145+
schedule: events.Schedule.rate(schedule),
146+
});
147+
rule.addTarget(new targets.LambdaFunction(this.fn));
148+
149+
NagSuppressions.addResourceSuppressions(this.fn, [
150+
{
151+
id: 'AwsSolutions-IAM4',
152+
reason: 'AWSLambdaBasicExecutionRole is required for CloudWatch Logs access',
153+
},
154+
{
155+
id: 'AwsSolutions-IAM5',
156+
reason: 'DynamoDB index/* wildcards generated by CDK grantReadWriteData/grantReadData for StatusIndex query access',
157+
},
158+
], true);
159+
}
160+
}

cdk/src/constructs/task-api.ts

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,15 @@ export interface TaskApiProps {
121121
*/
122122
readonly orchestratorFunctionArn?: string;
123123

124+
/**
125+
* Maximum concurrent tasks per user (#441). Threaded to the get-task
126+
* handler so the queue-wait ETA heuristic agrees with the
127+
* orchestrator's admission cap. Must match
128+
* ``TaskOrchestrator.maxConcurrentTasksPerUser``.
129+
* @default 10
130+
*/
131+
readonly maxConcurrentTasksPerUser?: number;
132+
124133
/**
125134
* API Gateway stage name.
126135
* @default 'v1'
@@ -533,7 +542,12 @@ export class TaskApi extends Construct {
533542
handler: 'handler',
534543
runtime: Runtime.NODEJS_24_X,
535544
architecture: Architecture.ARM_64,
536-
environment: commonEnv,
545+
environment: {
546+
...commonEnv,
547+
// #441: queue-position ETA heuristic must agree with the
548+
// orchestrator's per-user admission cap.
549+
MAX_CONCURRENT_TASKS_PER_USER: String(props.maxConcurrentTasksPerUser ?? 10),
550+
},
537551
bundling: commonBundling,
538552
});
539553

cdk/src/constructs/task-status.ts

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,23 @@
2121
* Valid task states in the task lifecycle state machine.
2222
*
2323
* States progress through the lifecycle:
24-
* [PENDING_UPLOADS ->] SUBMITTED -> HYDRATING ->
24+
* [PENDING_UPLOADS ->] SUBMITTED [<-> QUEUED] -> HYDRATING ->
2525
* RUNNING -> FINALIZING -> terminal (COMPLETED / FAILED / CANCELLED / TIMED_OUT).
2626
* See ORCHESTRATOR.md for the full state transition table.
2727
*
2828
* PENDING_UPLOADS is a pre-active state for tasks with presigned-upload
2929
* attachments: no compute allocated, no concurrency slot consumed. The
3030
* task transitions to SUBMITTED once uploads are confirmed and screened.
3131
*
32+
* QUEUED (#441) is a pre-active state for tasks that hit the per-user
33+
* admission cap: the admission slot was NOT acquired, so no compute is
34+
* allocated and no concurrency slot is consumed. The admission-queue
35+
* pickup Lambda re-attempts admission in FIFO order (by ``created_at``)
36+
* as slots free up, transitioning QUEUED -> SUBMITTED and re-invoking
37+
* the orchestrator. A pickup that loses the admission race simply
38+
* re-queues (SUBMITTED -> QUEUED); FIFO position is preserved because
39+
* ``created_at`` never changes.
40+
*
3241
* AWAITING_APPROVAL is the Cedar-HITL soft-deny gate surface: the
3342
* task is alive but paused on a human decision. See
3443
* `docs/design/CEDAR_HITL_GATES.md` §10.3 for the joint
@@ -37,6 +46,7 @@
3746
*/
3847
export const TaskStatus = {
3948
PENDING_UPLOADS: 'PENDING_UPLOADS',
49+
QUEUED: 'QUEUED',
4050
SUBMITTED: 'SUBMITTED',
4151
HYDRATING: 'HYDRATING',
4252
RUNNING: 'RUNNING',
@@ -66,10 +76,14 @@ export const TERMINAL_STATUSES: readonly TaskStatusType[] = [
6676
/**
6777
* Pre-active states where the task exists but has not entered the
6878
* orchestration pipeline. No compute resources are allocated and no
69-
* concurrency slot is consumed.
79+
* concurrency slot is consumed. QUEUED belongs here (#441): admission
80+
* explicitly did NOT acquire a slot, so counting it as active would
81+
* deadlock the queue (queued tasks would hold the very slots they
82+
* are waiting for).
7083
*/
7184
export const PRE_ACTIVE_STATUSES: readonly TaskStatusType[] = [
7285
TaskStatus.PENDING_UPLOADS,
86+
TaskStatus.QUEUED,
7387
];
7488

7589
/**
@@ -97,7 +111,11 @@ export const VALID_TRANSITIONS: Readonly<Record<TaskStatusType, readonly TaskSta
97111
// Transitions to SUBMITTED on confirm-uploads success, FAILED on
98112
// screening failure, CANCELLED on user cancel or 30-min auto-cancel.
99113
[TaskStatus.PENDING_UPLOADS]: [TaskStatus.SUBMITTED, TaskStatus.FAILED, TaskStatus.CANCELLED],
100-
[TaskStatus.SUBMITTED]: [TaskStatus.HYDRATING, TaskStatus.FAILED, TaskStatus.CANCELLED],
114+
// QUEUED (#441): admission-capped task awaiting a free concurrency
115+
// slot. SUBMITTED on pickup (slot acquired), CANCELLED on user
116+
// cancel, FAILED only via the queue-stranded backstop.
117+
[TaskStatus.QUEUED]: [TaskStatus.SUBMITTED, TaskStatus.CANCELLED, TaskStatus.FAILED],
118+
[TaskStatus.SUBMITTED]: [TaskStatus.QUEUED, TaskStatus.HYDRATING, TaskStatus.FAILED, TaskStatus.CANCELLED],
101119
[TaskStatus.HYDRATING]: [
102120
TaskStatus.RUNNING,
103121
TaskStatus.AWAITING_APPROVAL,

cdk/src/handlers/get-task.ts

Lines changed: 84 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,10 @@
1818
*/
1919

2020
import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
21-
import { DynamoDBDocumentClient, GetCommand } from '@aws-sdk/lib-dynamodb';
21+
import { DynamoDBDocumentClient, GetCommand, QueryCommand } from '@aws-sdk/lib-dynamodb';
2222
import type { APIGatewayProxyEvent, APIGatewayProxyResult } from 'aws-lambda';
2323
import { ulid } from 'ulid';
24+
import { TaskStatus } from '../constructs/task-status';
2425
import { extractUserId } from './shared/gateway';
2526
import { logger } from './shared/logger';
2627
import { ErrorCode, errorResponse, successResponse } from './shared/response';
@@ -29,6 +30,81 @@ import { type TaskRecord, toTaskDetail } from './shared/types';
2930
const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({}));
3031
const TABLE_NAME = process.env.TASK_TABLE_NAME!;
3132

33+
/** Must match the orchestrator's per-user admission cap for a sane ETA. */
34+
const MAX_CONCURRENT = Number(process.env.MAX_CONCURRENT_TASKS_PER_USER ?? '10');
35+
36+
/**
37+
* Rough per-task duration used for the queue-wait heuristic (#441).
38+
* Deliberately coarse — the ETA is a UX hint ("minutes, not seconds"),
39+
* not a promise. Override via env for fleets with unusual task profiles.
40+
*/
41+
const QUEUE_ETA_AVG_TASK_DURATION_S = Number(process.env.QUEUE_ETA_AVG_TASK_DURATION_S ?? '600');
42+
43+
/**
44+
* Compute the task's 1-based FIFO position among the user's QUEUED tasks
45+
* (#441) plus a coarse wait estimate. Position is per user because the
46+
* admission cap is per user — a global position would overstate the wait
47+
* for users whose slots are free.
48+
*
49+
* Queries the UserStatusIndex GSI (PK user_id, SK status_created_at
50+
* ``QUEUED#...``) and ranks by ``created_at``, which is the pickup
51+
* Lambda's FIFO key (``status_created_at`` moves on re-queue;
52+
* ``created_at`` never does).
53+
*
54+
* Fail-open: any error returns undefined so a GSI hiccup degrades the
55+
* response to ``queue_position: null`` instead of failing the whole GET.
56+
*/
57+
async function computeQueueInfo(
58+
record: TaskRecord,
59+
): Promise<{ queue_position: number; estimated_wait_s: number | null } | undefined> {
60+
try {
61+
let ahead = 0;
62+
let found = false;
63+
let lastKey: Record<string, unknown> | undefined;
64+
do {
65+
const resp = await ddb.send(new QueryCommand({
66+
TableName: TABLE_NAME,
67+
IndexName: 'UserStatusIndex',
68+
KeyConditionExpression: 'user_id = :uid AND begins_with(status_created_at, :queued)',
69+
ExpressionAttributeValues: {
70+
':uid': record.user_id,
71+
':queued': `${TaskStatus.QUEUED}#`,
72+
},
73+
ProjectionExpression: 'task_id, created_at',
74+
ExclusiveStartKey: lastKey as Record<string, never> | undefined,
75+
}));
76+
for (const item of resp.Items ?? []) {
77+
if (item.task_id === record.task_id) {
78+
found = true;
79+
} else if (typeof item.created_at === 'string' && item.created_at < record.created_at) {
80+
ahead++;
81+
}
82+
}
83+
lastKey = resp.LastEvaluatedKey;
84+
} while (lastKey);
85+
86+
// The task left QUEUED between our GetItem and this query — report
87+
// no queue info rather than a stale position.
88+
if (!found) {
89+
return undefined;
90+
}
91+
92+
const position = ahead + 1;
93+
// Coarse ETA: slots drain MAX_CONCURRENT at a time, each batch
94+
// lasting roughly one average task duration.
95+
const estimatedWaitS = MAX_CONCURRENT > 0 && QUEUE_ETA_AVG_TASK_DURATION_S > 0
96+
? Math.ceil(position / MAX_CONCURRENT) * QUEUE_ETA_AVG_TASK_DURATION_S
97+
: null;
98+
return { queue_position: position, estimated_wait_s: estimatedWaitS };
99+
} catch (err) {
100+
logger.warn('Queue position lookup failed (fail-open — returning null position)', {
101+
task_id: record.task_id,
102+
error: err instanceof Error ? err.message : String(err),
103+
});
104+
return undefined; // nosemgrep: ts-silent-success-masking -- queue position is a best-effort UX hint; a GSI failure must not fail the whole GET /tasks/{id}
105+
}
106+
}
107+
32108
/**
33109
* GET /v1/tasks/{task_id} — Get full task details.
34110
*/
@@ -64,8 +140,13 @@ export async function handler(event: APIGatewayProxyEvent): Promise<APIGatewayPr
64140
return errorResponse(403, ErrorCode.FORBIDDEN, 'You do not have access to this task.', requestId);
65141
}
66142

67-
// 5. Return task detail
68-
return successResponse(200, toTaskDetail(record), requestId);
143+
// 5. For QUEUED tasks, compute read-time queue position + ETA (#441)
144+
const queueInfo = record.status === TaskStatus.QUEUED
145+
? await computeQueueInfo(record)
146+
: undefined;
147+
148+
// 6. Return task detail
149+
return successResponse(200, toTaskDetail(record, queueInfo), requestId);
69150
} catch (err) {
70151
logger.error('Failed to get task', { error: String(err), request_id: requestId });
71152
return errorResponse(500, ErrorCode.INTERNAL_ERROR, 'Internal server error.', requestId);

0 commit comments

Comments
 (0)